added from_sql and to_sql methods to MsgData

msg.py

@@ -1,7 +1,9 @@
 import asyncio
 import logging
 import os
+import sqlite3
 from datetime import datetime, timedelta
+from pathlib import Path
 from typing import Dict, Iterable
 
 import discord
@@ -26,6 +28,40 @@ class MsgData:
         self.lock = asyncio.Lock()
         return self
 
+    @classmethod
+    def from_sql(cls, db):
+        if isinstance(db, (str, Path)):
+            con = sqlite3.connect(db)
+        elif isinstance(db, sqlite3.Connection):
+            con = db
+
+        self = MsgData()
+        self.msgs: pd.DataFrame = pd.read_sql('select * from msgs', con=con, index_col='id')
+        self.msgs['created'] = self.msgs['created'].apply(pd.to_datetime, utc=True)
+        self.reactions: pd.DataFrame = pd.read_sql('select * from reactions', con).set_index(['msg id', 'emoji'])
+        return self
+
+    def to_sql(self, db):
+        if isinstance(db, (str, Path)):
+            con = sqlite3.connect(db)
+        elif isinstance(db, sqlite3.Connection):
+            con = db
+
+        self.msgs.drop('object', axis=1).to_sql(
+            name='msgs',
+            con=con,
+            if_exists='replace',
+            index=True,
+            index_label=self.msgs.index.name
+        )
+        self.reactions.to_sql(
+            name='reactions',
+            con=con,
+            if_exists='replace',
+            index=True,
+            index_label=self.reactions.index.name
+        )
+
     def __str__(self):
         return str(self.msgs) + '\n\n' + str(self.reactions)
 
@@ -63,9 +99,21 @@ class MsgData:
     def cancellations(self, days: int = 14):
         return cancellations(msg_df=self.msgs, react_df=self.reactions, days=days)
 
-    def cancellation_totals(self, days: int = 14):
+    def cancellation_totals(self, days):
         return cancelled_totals(cdf=self.cancellations(days=days))
 
+    def cancellation_leaderboard(self, days, top: int = None):
+        df = self.cancellation_totals(days)
+        if top is not None:
+            df = df.iloc[:top]
+        width = max(list(map(lambda s: len(str(s)), df.index.values)))
+        res = f'Cancellation totals, past {days} days\n'
+        res += '\n'.join(
+            f"`{name.ljust(width + 1)}with {row['total']:<2.0f} total`"
+            for name, row in df.iterrows()
+        )
+        return res
+
 
 async def message_df(client: discord.Client, **kwargs):
     return pd.DataFrame(
@@ -94,6 +142,7 @@ def message_dict(m: discord.Message) -> Dict:
         'id': m.id,
         'created': m.created_at.astimezone(),
         'display_name': m.author.display_name,
+        'user id': m.author.id,
         'message': m.content,
         'channel': m.channel.name,
         'channel link': f'<#{m.channel.id}>',
@@ -106,13 +155,11 @@ async def reaction_df(msgs: Iterable[discord.Message]):
 
 
 async def reaction_series(msg: discord.Message):
-    return pd.DataFrame(
-        [
-            await reaction_dict(r)
-            for r in msg.reactions
-            if isinstance(r.emoji, discord.Emoji)
-        ]
-    )
+    return pd.DataFrame([
+        await reaction_dict(r)
+        for r in msg.reactions
+        if isinstance(r.emoji, discord.Emoji)
+    ])
 
 
 async def reaction_dict(r: discord.Reaction) -> Dict:
@@ -160,14 +207,6 @@ def cancelled_totals(cdf: pd.DataFrame) -> pd.DataFrame:
     }).sort_values('total', ascending=False)
 
 
-def report_string(df):
-    width = max(list(map(lambda s: len(str(s)), df.index.values)))
-    return '\n'.join(
-        f"`{name.ljust(width + 1)}with {row['total']:<2.0f} total`"
-        for name, row in df.iterrows()
-    )
-
-
 if __name__ == '__main__':
     client = discord.Client()
 

robopage.py

@@ -6,8 +6,8 @@ from threading import Lock
 import discord
 from dotenv import load_dotenv
 
+import jokes
 import msg
-from jokes import CumJoke, BlackJoke, AssJoke, DominosJoke
 
 logging.basicConfig(level=logging.INFO)
 
@@ -21,12 +21,9 @@ class RoboPage(discord.Client):
 
     def __init__(self, *args, **kwargs):
         super(RoboPage, self).__init__(*args, **kwargs)
-        self.jokes = [
-            CumJoke(),
-            BlackJoke(),
-            AssJoke(),
-            DominosJoke()
-        ]
+        attrs = filter(lambda n: n.endswith('Joke') and not n.startswith('Joke'), dir(jokes))
+        attrs = map(lambda n: getattr(jokes, n)(), attrs)
+        self.jokes = list(attrs)
         self.lock = Lock()
 
     def run(self):
@@ -40,9 +37,9 @@ class RoboPage(discord.Client):
         self.data: msg.MsgData = await msg.MsgData.create(
             client=self,
             limit=3000,
-            # limit=20,
             days=14,
         )
+        self.data.to_sql('messages.db')
         LOGGER.info(str(self.data.msgs.columns))
         LOGGER.info(str(self.data.reactions.columns))
 
@@ -51,7 +48,7 @@ class RoboPage(discord.Client):
 
         if message.author != self.user:
             if 'most cancelled' in message.content:
-                await message.reply(self.get_cancelled_totals(days=14))
+                await message.reply(self.data.cancellation_leaderboard(14))
 
             elif (m := re.search('top cancelled (?P<name>\w+)', message.content)) is not None:
                 async with self.data.lock:
@@ -62,11 +59,6 @@ class RoboPage(discord.Client):
                     print(f'{joke.__class__.__name__} detected:\n{message.content}\n{scan_res}')
                     await joke.respond(message, self, scan_res)
 
-    def get_cancelled_totals(self, days):
-        res = self.data.cancellation_totals(days)
-        res = f'Cancellation totals, past {days} days\n' + msg.report_string(res.iloc[:5])
-        return res
-
     def top_cancellations(self, user: str, days: int):
         cdf = self.data.cancellations(days)
         cdf = cdf[cdf['display_name'].str.contains(user, case=False)]
@@ -88,7 +80,7 @@ if __name__ == '__main__':
 
     @client.event
     async def on_ready():
-        print(f'{client.user} has connected to Discord!')
+        print(len(list(client.get_all_members())))
         await client.handle_ready()
         print(client.data.cancellation_totals(14))