made ReactionData

kwaylon/kwaylon.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 import re
 from datetime import timedelta, datetime
@@ -9,7 +10,7 @@ from nextcord import Client, Message, TextChannel
 from nextcord import RawReactionActionEvent
 
 from . import jokes
-from .data import MsgData
+from .reactions import ReactionData
 
 LIL_STINKY_ID = 704043422276780072
 
@@ -23,6 +24,7 @@ class Kwaylon(Client):
         super().__init__(*args, **kwargs)
         self.limit, self.days = limit, days
         self.jokes = list(jokes.collect_jokes())
+        self.lock = asyncio.Lock()
 
         self.most_regex = re.compile('most\s+(?P<emoji>\S+)')
         self.leaderboard_regex = re.compile(
@@ -38,11 +40,9 @@ class Kwaylon(Client):
 
         # await alive()
 
-        self.data = MsgData(self.db_path)
+        self.data = ReactionData(self.db_path)
-
+        # await self.data.scan_messages(client=self, limit=100)
-        await self.data.load_sql()
+        self.data.read_all()
-        if not hasattr(self.data, 'reactions'):
-            await self.data.scan_messages(client=self, limit=self.limit, days=self.days)
 
     async def handle_message(self, message: Message):
         if message.author != self.user:
@@ -57,11 +57,15 @@ class Kwaylon(Client):
                     return
 
             if (m := self.most_regex.match(message.clean_content)) is not None:
-                await self.data.load_sql()
+                # await self.data.load_sql()
                 emoji = get_emoji_name(m.group('emoji'))
                 LOGGER.info(emoji)
 
-                if (df := await self.data.get_emoji_info(emoji)) is not None and df.shape[0] > 0:
+                with self.data.connect() as con:
+                    df = self.data.read_emoji(emoji, con)
+                con.close()
+
+                if df.shape[0] > 0:
                     kwargs = {'days': 14}
                     if (day_match := re.search('(?P<days>\d+) days', message.content)):
                         days = int(day_match.group('days'))
@@ -70,8 +74,9 @@ class Kwaylon(Client):
                     if 'leaderboard' in message.content:
                         LOGGER.info(f'Building leaderboard')
                         res = f'{m.group("emoji")} totals, past {kwargs["days"]} days\n'
-                        res += await self.leaderboard(df, **kwargs)
+                        if (board := await self.leaderboard(df, **kwargs)) is not None:
-                        await message.reply(res)
+                            res += board
+                            await message.reply(res)
                         LOGGER.info(f'Done')
                     else:
                         LOGGER.info(f'Most {m.group("emoji")}')
@@ -91,36 +96,30 @@ class Kwaylon(Client):
         valid_dates = df['datetime'] > start
         df = df.loc[valid_dates]
 
-        df = df.groupby('auth_id').sum()
+        if df.shape[0] > 0:
-        counts = df['count'].sort_values(ascending=False)
+            df = df.groupby('auth_id').sum()
-        counts.index = [(await self.fetch_user(idx)).display_name for idx in counts.index]
+            counts = df['count'].sort_values(ascending=False)
+            counts.index = [(await self.fetch_user(idx)).display_name for idx in counts.index]
 
-        width = max([len(str(s)) for s in counts.index])
+            width = max([len(str(s)) for s in counts.index])
 
-        res = '\n'.join(
+            res = '\n'.join(
-            f"`{str(name).ljust(width + 1)}with {cnt:<2.0f} total`"
+                f"`{str(name).ljust(width + 1)}with {cnt:<2.0f} total`"
-            for name, cnt in counts.iteritems()
+                for name, cnt in counts.iteritems()
-        )
+            )
-        return res
+            return res
 
     async def handle_raw_reaction(self, payload: RawReactionActionEvent):
         LOGGER.info(payload)
+
         guild = await self.fetch_guild(payload.guild_id)
         channel = await guild.fetch_channel(payload.channel_id)
         message = await channel.fetch_message(payload.message_id)
 
-        if payload.event_type == 'REACTION_ADD':
+        with self.lock:
-            LOGGER.info(
+            with self.data.connect() as con:
-                f'{payload.member.display_name} added {payload.emoji} to\n' + \
+                self.data.add_reactions_from_message(message, con)
-                f'{message.author.display_name}: {message.content}')
+            con.close()
-            # await self.data.add_reaction()
-
-        elif payload.event_type == 'REACTION_REMOVE':
-            LOGGER.info(f'{payload.emoji} removed from\n{message.author}: {message.content}')
-            # await self.data.remove_reaction(event=payload)
-
-        if hasattr(self, 'data'):
-            await self.data.update_reaction(msg=message)
 
     async def fetch_message(self, row: pd.Series):
         guild = await self.fetch_guild(row['guild_id'])

kwaylon/reactions.py

@@ -0,0 +1,60 @@
+import logging
+import sqlite3
+from dataclasses import dataclass
+from pathlib import Path
+
+import pandas as pd
+from nextcord import Message, Client
+
+from .msg import reaction_dict, message_gen
+
+LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class ReactionData:
+    path: Path
+
+    def connect(self, *args, **kwargs) -> sqlite3.Connection:
+        return sqlite3.connect(self.path, *args, **kwargs)
+
+    async def scan_messages(self, client: Client, **kwargs):
+        try:
+            with self.connect() as con:
+                async for msg in message_gen(client=client, **kwargs):
+                    if len(msg.reactions) > 0:
+                        self.add_reactions_from_message(msg, con)
+        except Exception as e:
+            LOGGER.exception(e)
+        finally:
+            con.close()
+
+    def add_reactions_from_message(self, msg: Message, con: sqlite3.Connection = None):
+        con = con or sqlite3.connect(self.path)
+
+        try:
+            con.execute(f'DELETE FROM reactions WHERE msg_id = {msg.id}')
+
+            data = [tuple(reaction_dict(reaction).values()) for reaction in msg.reactions]
+            query = f'INSERT INTO reactions VALUES({",".join("?" for _ in range(8))})'
+            con.executemany(query, data)
+        except Exception as e:
+            LOGGER.exception(e)
+        # else:
+        #     LOGGER.info(f'Wrote {len(data)} rows to {self.path.name}')
+
+    def read_emoji(self, emoji: str, con: sqlite3.Connection = None) -> pd.DataFrame:
+        return self.read_sql(query=f"SELECT * FROM reactions WHERE emoji LIKE '{emoji}'", con=con)
+
+    def read_all(self, con: sqlite3.Connection = None) -> pd.DataFrame:
+        return self.read_sql(query='SELECT * FROM reactions', con=con)
+
+    def read_sql(self, query: str, con: sqlite3.Connection = None):
+        con = con or sqlite3.connect(self.path)
+
+        res = pd.read_sql(query, con=con, index_col=None)
+        LOGGER.info(f'Read {res.shape[0]} reactions')
+
+        res['datetime'] = pd.to_datetime(res['datetime'])
+
+        return res.sort_values('count', ascending=False)