clean role mentions
This commit is contained in:
parent
648a50aded
commit
ea87aa1f1f
3 changed files with 91 additions and 28 deletions
|
|
@ -12,6 +12,7 @@ class DiscordReader:
|
||||||
|
|
||||||
self.guild: discord.Guild | None = None
|
self.guild: discord.Guild | None = None
|
||||||
self.client: discord.Client | None = None
|
self.client: discord.Client | None = None
|
||||||
|
self.role_map: Dict[int, str] = {}
|
||||||
|
|
||||||
def _create_client(self):
|
def _create_client(self):
|
||||||
intents = discord.Intents.default()
|
intents = discord.Intents.default()
|
||||||
|
|
@ -25,10 +26,16 @@ class DiscordReader:
|
||||||
if not self.client or self.client.is_closed():
|
if not self.client or self.client.is_closed():
|
||||||
self.client = self._create_client()
|
self.client = self._create_client()
|
||||||
|
|
||||||
|
# login() initializes the internal HTTP session needed for API calls
|
||||||
await self.client.login(self.token)
|
await self.client.login(self.token)
|
||||||
# In a real app we'd wait until ready, here we simulate fetching the guild
|
|
||||||
|
# Use fetch methods specifically to bypass dependency on gateway cache
|
||||||
self.guild = await self.client.fetch_guild(self.server_id)
|
self.guild = await self.client.fetch_guild(self.server_id)
|
||||||
|
|
||||||
|
# Pre-fetch roles via API
|
||||||
|
roles = await self.guild.fetch_roles()
|
||||||
|
self.role_map = {r.id: r.name for r in roles}
|
||||||
|
|
||||||
async def validate(self) -> Dict[str, Any]:
|
async def validate(self) -> Dict[str, Any]:
|
||||||
"""Validates the token, server ID, intents, and permissions."""
|
"""Validates the token, server ID, intents, and permissions."""
|
||||||
results = {
|
results = {
|
||||||
|
|
|
||||||
|
|
@ -8,28 +8,50 @@ from src.core.base import MigrationContext
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
def clean_mentions(content: str, guild, mentions=None) -> str:
|
def clean_mentions(content: str, guild, user_mentions=None, role_mentions=None, role_map=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
|
# 1. Try provided guild
|
||||||
member = guild.get_member(uid)
|
member = guild.get_member(uid)
|
||||||
if member:
|
if member:
|
||||||
return f"`@{member.display_name}`"
|
return f"`@{member.display_name}`"
|
||||||
# Try message mentions list
|
# 2. Try message's user_mentions
|
||||||
if mentions:
|
if user_mentions:
|
||||||
for user in mentions:
|
for u in user_mentions:
|
||||||
if user.id == uid:
|
if u.id == uid:
|
||||||
name = getattr(user, 'display_name', user.name)
|
return f"`@{getattr(u, 'display_name', u.name)}`"
|
||||||
return f"`@{name}`"
|
# 3. Try global cache via guild.client
|
||||||
|
if hasattr(guild, 'client'):
|
||||||
|
user = guild.client.get_user(uid)
|
||||||
|
if user:
|
||||||
|
return f"`@{user.name}`"
|
||||||
return match.group(0)
|
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)
|
# 1. Try provided guild cache/list
|
||||||
return f"`@{role.name}`" if role else match.group(0)
|
role = guild.get_role(rid) or discord.utils.get(guild.roles, id=rid)
|
||||||
|
# 2. Try message's role_mentions
|
||||||
|
if not role and role_mentions:
|
||||||
|
role = discord.utils.get(role_mentions, id=rid)
|
||||||
|
|
||||||
|
# 3. Try all guilds the client is aware of (fallback for cache issues)
|
||||||
|
if not role and hasattr(guild, 'client'):
|
||||||
|
for g in guild.client.guilds:
|
||||||
|
role = g.get_role(rid)
|
||||||
|
if role: break
|
||||||
|
|
||||||
|
if role and role.name:
|
||||||
|
return f"`@{role.name}`"
|
||||||
|
|
||||||
|
# 4. Try provided role map
|
||||||
|
if role_map and rid in role_map:
|
||||||
|
return f"`@{role_map[rid]}`"
|
||||||
|
|
||||||
|
return match.group(0)
|
||||||
|
|
||||||
def replace_channel(match):
|
def replace_channel(match):
|
||||||
cid = int(match.group(1))
|
cid = int(match.group(1))
|
||||||
|
|
@ -121,8 +143,8 @@ async def migrate_messages(context: MigrationContext, source_channel_id: int, ta
|
||||||
)
|
)
|
||||||
continue
|
continue
|
||||||
else:
|
else:
|
||||||
# Use custom clean_mentions with msg.mentions for accuracy
|
# Use custom clean_mentions with msg mentions for accuracy
|
||||||
content = clean_mentions(msg.content, msg.guild, msg.mentions)
|
content = clean_mentions(msg.content, msg.guild, msg.mentions, msg.role_mentions, context.discord_reader.role_map)
|
||||||
|
|
||||||
# Process attachments
|
# Process attachments
|
||||||
files = []
|
files = []
|
||||||
|
|
@ -144,7 +166,13 @@ 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, snapshot.mentions if hasattr(snapshot, 'mentions') else None)
|
content = clean_mentions(
|
||||||
|
content,
|
||||||
|
msg.guild,
|
||||||
|
snapshot.mentions if hasattr(snapshot, 'mentions') else None,
|
||||||
|
snapshot.role_mentions if hasattr(snapshot, 'role_mentions') else None,
|
||||||
|
context.discord_reader.role_map
|
||||||
|
)
|
||||||
# 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")
|
||||||
|
|
|
||||||
|
|
@ -8,28 +8,50 @@ from src.core.base import MigrationContext
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
def clean_mentions(content: str, guild, mentions=None) -> str:
|
def clean_mentions(content: str, guild, user_mentions=None, role_mentions=None, role_map=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
|
# 1. Try provided guild
|
||||||
member = guild.get_member(uid)
|
member = guild.get_member(uid)
|
||||||
if member:
|
if member:
|
||||||
return f"`@{member.display_name}`"
|
return f"`@{member.display_name}`"
|
||||||
# Try message mentions list
|
# 2. Try message's user_mentions
|
||||||
if mentions:
|
if user_mentions:
|
||||||
for user in mentions:
|
for u in user_mentions:
|
||||||
if user.id == uid:
|
if u.id == uid:
|
||||||
name = getattr(user, 'display_name', user.name)
|
return f"`@{getattr(u, 'display_name', u.name)}`"
|
||||||
return f"`@{name}`"
|
# 3. Try global cache via guild.client
|
||||||
|
if hasattr(guild, 'client'):
|
||||||
|
user = guild.client.get_user(uid)
|
||||||
|
if user:
|
||||||
|
return f"`@{user.name}`"
|
||||||
return match.group(0)
|
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)
|
# 1. Try provided guild cache/list
|
||||||
return f"`@{role.name}`" if role else match.group(0)
|
role = guild.get_role(rid) or discord.utils.get(guild.roles, id=rid)
|
||||||
|
# 2. Try message's role_mentions
|
||||||
|
if not role and role_mentions:
|
||||||
|
role = discord.utils.get(role_mentions, id=rid)
|
||||||
|
|
||||||
|
# 3. Try all guilds the client is aware of (fallback for cache issues)
|
||||||
|
if not role and hasattr(guild, 'client'):
|
||||||
|
for g in guild.client.guilds:
|
||||||
|
role = g.get_role(rid)
|
||||||
|
if role: break
|
||||||
|
|
||||||
|
if role and role.name:
|
||||||
|
return f"`@{role.name}`"
|
||||||
|
|
||||||
|
# 4. Try provided role map
|
||||||
|
if role_map and rid in role_map:
|
||||||
|
return f"`@{role_map[rid]}`"
|
||||||
|
|
||||||
|
return match.group(0)
|
||||||
|
|
||||||
def replace_channel(match):
|
def replace_channel(match):
|
||||||
cid = int(match.group(1))
|
cid = int(match.group(1))
|
||||||
|
|
@ -123,8 +145,8 @@ async def migrate_messages(context: MigrationContext, source_channel_id: int, ta
|
||||||
)
|
)
|
||||||
continue
|
continue
|
||||||
else:
|
else:
|
||||||
# Use custom clean_mentions with msg.mentions for accuracy
|
# Use custom clean_mentions with msg mentions for accuracy
|
||||||
content = clean_mentions(msg.content, msg.guild, msg.mentions)
|
content = clean_mentions(msg.content, msg.guild, msg.mentions, msg.role_mentions, context.discord_reader.role_map)
|
||||||
|
|
||||||
# Process attachments
|
# Process attachments
|
||||||
files = []
|
files = []
|
||||||
|
|
@ -142,7 +164,13 @@ 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, snapshot.mentions if hasattr(snapshot, 'mentions') else None)
|
content = clean_mentions(
|
||||||
|
content,
|
||||||
|
msg.guild,
|
||||||
|
snapshot.mentions if hasattr(snapshot, 'mentions') else None,
|
||||||
|
snapshot.role_mentions if hasattr(snapshot, 'role_mentions') else None,
|
||||||
|
context.discord_reader.role_map
|
||||||
|
)
|
||||||
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")
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue