clean user mentions

This commit is contained in:
rambros 2026-02-27 03:15:13 +05:30
parent 0d1e0f6717
commit 648a50aded
2 changed files with 32 additions and 12 deletions

View file

@ -8,19 +8,28 @@ from src.core.base import MigrationContext
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def clean_mentions(content: str, guild) -> str: def clean_mentions(content: str, guild, mentions=None) -> str:
if not content or not guild: if not content or not guild:
return content return content
def replace_user(match): def replace_user(match):
uid = int(match.group(1)) uid = int(match.group(1))
# Try guild cache first
member = guild.get_member(uid) member = guild.get_member(uid)
return f"@{member.display_name}" if member else match.group(0) if member:
return f"`@{member.display_name}`"
# Try message mentions list
if mentions:
for user in mentions:
if user.id == uid:
name = getattr(user, 'display_name', user.name)
return f"`@{name}`"
return match.group(0)
def replace_role(match): def replace_role(match):
rid = int(match.group(1)) rid = int(match.group(1))
role = guild.get_role(rid) role = guild.get_role(rid)
return f"@{role.name}" if role else match.group(0) return f"`@{role.name}`" if role else match.group(0)
def replace_channel(match): def replace_channel(match):
cid = int(match.group(1)) cid = int(match.group(1))
@ -30,6 +39,7 @@ def clean_mentions(content: str, guild) -> str:
content = re.sub(r'<@!?([0-9]+)>', replace_user, content) content = re.sub(r'<@!?([0-9]+)>', replace_user, content)
content = re.sub(r'<@&([0-9]+)>', replace_role, content) content = re.sub(r'<@&([0-9]+)>', replace_role, content)
content = re.sub(r'<#([0-9]+)>', replace_channel, content) content = re.sub(r'<#([0-9]+)>', replace_channel, content)
content = content.replace("@everyone", "`@everyone`").replace("@here", "`@here`")
return content return content
@ -111,8 +121,8 @@ async def migrate_messages(context: MigrationContext, source_channel_id: int, ta
) )
continue continue
else: else:
# Get clean content # Use custom clean_mentions with msg.mentions for accuracy
content = msg.clean_content content = clean_mentions(msg.content, msg.guild, msg.mentions)
# Process attachments # Process attachments
files = [] files = []
@ -134,7 +144,7 @@ async def migrate_messages(context: MigrationContext, source_channel_id: int, ta
if not content: # Only update content if it wasn't already set (e.g., by thread_starter_message) if not content: # Only update content if it wasn't already set (e.g., by thread_starter_message)
content = snapshot.content content = snapshot.content
if hasattr(msg, 'guild') and msg.guild: if hasattr(msg, 'guild') and msg.guild:
content = clean_mentions(content, msg.guild) content = clean_mentions(content, msg.guild, snapshot.mentions if hasattr(snapshot, 'mentions') else None)
# Add snapshot attachments to the list to process # Add snapshot attachments to the list to process
attachments_to_process.extend(snapshot.attachments) attachments_to_process.extend(snapshot.attachments)
logger.debug(f"Found forwarded snapshot content: {content[:50]}... and {len(snapshot.attachments)} attachments") logger.debug(f"Found forwarded snapshot content: {content[:50]}... and {len(snapshot.attachments)} attachments")

View file

@ -8,19 +8,28 @@ from src.core.base import MigrationContext
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def clean_mentions(content: str, guild) -> str: def clean_mentions(content: str, guild, mentions=None) -> str:
if not content or not guild: if not content or not guild:
return content return content
def replace_user(match): def replace_user(match):
uid = int(match.group(1)) uid = int(match.group(1))
# Try guild cache first
member = guild.get_member(uid) member = guild.get_member(uid)
return f"@{member.display_name}" if member else match.group(0) if member:
return f"`@{member.display_name}`"
# Try message mentions list
if mentions:
for user in mentions:
if user.id == uid:
name = getattr(user, 'display_name', user.name)
return f"`@{name}`"
return match.group(0)
def replace_role(match): def replace_role(match):
rid = int(match.group(1)) rid = int(match.group(1))
role = guild.get_role(rid) role = guild.get_role(rid)
return f"@{role.name}" if role else match.group(0) return f"`@{role.name}`" if role else match.group(0)
def replace_channel(match): def replace_channel(match):
cid = int(match.group(1)) cid = int(match.group(1))
@ -30,6 +39,7 @@ def clean_mentions(content: str, guild) -> str:
content = re.sub(r'<@!?([0-9]+)>', replace_user, content) content = re.sub(r'<@!?([0-9]+)>', replace_user, content)
content = re.sub(r'<@&([0-9]+)>', replace_role, content) content = re.sub(r'<@&([0-9]+)>', replace_role, content)
content = re.sub(r'<#([0-9]+)>', replace_channel, content) content = re.sub(r'<#([0-9]+)>', replace_channel, content)
content = content.replace("@everyone", "`@everyone`").replace("@here", "`@here`")
return content return content
@ -113,8 +123,8 @@ async def migrate_messages(context: MigrationContext, source_channel_id: int, ta
) )
continue continue
else: else:
# Get clean content # Use custom clean_mentions with msg.mentions for accuracy
content = msg.clean_content content = clean_mentions(msg.content, msg.guild, msg.mentions)
# Process attachments # Process attachments
files = [] files = []
@ -132,7 +142,7 @@ async def migrate_messages(context: MigrationContext, source_channel_id: int, ta
if not content: if not content:
content = snapshot.content content = snapshot.content
if hasattr(msg, 'guild') and msg.guild: if hasattr(msg, 'guild') and msg.guild:
content = clean_mentions(content, msg.guild) content = clean_mentions(content, msg.guild, snapshot.mentions if hasattr(snapshot, 'mentions') else None)
attachments_to_process.extend(snapshot.attachments) attachments_to_process.extend(snapshot.attachments)
logger.debug(f"Found forwarded snapshot content: {content[:50]}... and {len(snapshot.attachments)} attachments") logger.debug(f"Found forwarded snapshot content: {content[:50]}... and {len(snapshot.attachments)} attachments")