updated message_gen for nextcord

kwaylon/msg.py

@@ -2,34 +2,46 @@ import logging
 from datetime import datetime, timedelta
 from typing import Dict, Iterable
 
-import discord
 import pandas as pd
+from nextcord import Client, Message
+from nextcord import Reaction
+from nextcord import TextChannel
+from nextcord.utils import AsyncIterator
 
 LOGGER = logging.getLogger(__name__)
 
 
-async def message_df(client: discord.Client, **kwargs):
+async def message_df(client: Client, **kwargs):
     return pd.DataFrame(
         [message_dict(m) async for m in message_gen(client, **kwargs)]
     ).set_index('id').sort_values('created', ascending=False)
 
 
-async def message_gen(client: discord.Client, limit=20, days: int = 90, **kwargs):
+async def message_gen(client: Client, limit=20, days: int = 90, **kwargs) -> AsyncIterator:
-    channels = client.get_all_channels()
-    channels = filter(lambda c: isinstance(c, discord.TextChannel), channels)
-    channels = filter(lambda c: c.category.name != 'Archive', channels)
-    channels = sorted(channels, key=lambda c: (c.category.name, c.name))
-    for channel in channels:
-        LOGGER.info(f'{channel.category.name} #{channel.name}')
     if 'after' not in kwargs:
         kwargs['after'] = (datetime.today() - timedelta(days=days))
     elif isinstance((after := kwargs.get('after', None)), datetime):
         kwargs['after'] = after.replace(tzinfo=None)
-        async for msg in channel.history(limit=limit, **kwargs):
+
+    kwargs['limit'] = limit
+
+    for channel in client.get_all_channels():
+        if channel.category is not None and channel.category.name != 'Archive':
+            if isinstance(channel, TextChannel):
+                LOGGER.info(f'Channel: {channel.category}: {channel.name}')
+                async for msg in channel.history(**kwargs):
                     yield msg
+                for thread in channel.threads:
+                    LOGGER.info(f'Thread: {channel.category}: {channel.name}: {thread.name}')
+                    async for msg in thread.history(**kwargs):
+                        yield msg
+            else:
+                continue
+    else:
+        LOGGER.info(f'Done getting messages')
 
 
-def message_dict(m: discord.Message) -> Dict:
+def message_dict(m: Message) -> Dict:
     return {
         'object': m,
         'id': m.id,
@@ -43,16 +55,16 @@ def message_dict(m: discord.Message) -> Dict:
     }
 
 
-def full_reaction_df(msgs: Iterable[discord.Message]):
+def full_reaction_df(msgs: Iterable[Message]):
     return pd.concat([reaction_df(msg) for msg in msgs])
 
 
-def reaction_df(msg: discord.Message):
+def reaction_df(msg: Message):
     df = pd.DataFrame([reaction_dict(r) for r in msg.reactions])
     return df.set_index(['msg id', 'emoji']) if not df.empty else df
 
 
-def reaction_dict(r: discord.Reaction) -> Dict:
+def reaction_dict(r: Reaction) -> Dict:
     return {
         'object': r,
         'msg id': r.message.id,