small improvements to loading reactions

src/kwaylon/kwaylon.py

@@ -2,7 +2,6 @@ import asyncio
 import logging
 import re
 import sqlite3
-from datetime import timedelta, datetime
 from pathlib import Path
 from typing import List
 
@@ -12,7 +11,7 @@ from nextcord import RawReactionActionEvent, Emoji
 from nextcord import utils
 
 from . import jokes
-from .reactions import ReactionData
+from .reactions import ReactionData, filter_days
 
 LIL_STINKY_ID = 704043422276780072
 
@@ -20,11 +19,12 @@ LOGGER = logging.getLogger(__name__)
 
 
 class Kwaylon(Client):
-    # db_path: Path = Path(r'../data/messages.db')
+    def __init__(self, limit: int = 5000, days: int = 30, db_path: Path = None, *args, **kwargs):
-
-    def __init__(self, limit: int = 5000, days: int = 30, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.db_path = Path.cwd() / 'data' / 'messages.db'
+        if db_path is None:
+            self.db_path = Path.cwd() / 'data' / 'messages.db'
+        else:
+            self.db_path = db_path
 
         self.limit, self.days = limit, days
         self.jokes = list(jokes.collect_jokes())
@@ -57,7 +57,7 @@ class Kwaylon(Client):
 
         try:
             self.data = ReactionData(self.db_path)
-            self.data.read_all()
+            LOGGER.info(f'{self.data.row_count():d} reactions in {self.db_path}')
         except sqlite3.Error as e:
             LOGGER.exception(e)
             LOGGER.error(f'self.db_path: {self.db_path}')
@@ -80,32 +80,33 @@ class Kwaylon(Client):
             emoji_name = get_emoji_name(emoji_ref)
             LOGGER.info(f'Most {emoji_name}')
 
-            with self.data.connect() as con:
+            async with message.channel.typing():
-                df = self.data.read_emoji(emoji_name, con)
+                with self.data.connect() as con:
-            con.close()
+                    df = self.data.read_emoji(emoji_name, con)
+                con.close()
 
-            days = get_days(message.content) or 14
+                days = get_days(message.content) or 14
-            df = filter_days(df, days)
+                df = filter_days(df, days)
 
-            if df.shape[0] > 0:
+                if df.shape[0] > 0:
-                LOGGER.info(f'{df.shape[0]} messages with {emoji_ref} after filtering')
+                    LOGGER.info(f'{df.shape[0]} messages with {emoji_ref} after filtering')
 
-                if 'leaderboard' in message.content:
+                    if 'leaderboard' in message.content:
-                    LOGGER.info(f'Building leaderboard')
+                        LOGGER.info(f'Building leaderboard')
-                    res = f'{emoji_ref} totals, past {days} days\n'
+                        res = f'{emoji_ref} totals, past {days} days\n'
-                    if (board := await self.leaderboard(df)) is not None:
+                        if (board := await self.leaderboard(df)) is not None:
-                        res += board
+                            res += board
-                        await message.reply(res)
+                            await message.reply(res)
+
+                    else:
+                        most = df.sort_values('count').iloc[-1]
+                        msg = await self.fetch_message(most)
+                        await message.reply(f'{msg.jump_url}')
 
                 else:
-                    most = df.sort_values('count').iloc[-1]
+                    await message.reply(f"NObody (in the past {days} days)...gah, leave me alone!")
-                    msg = await self.fetch_message(most)
-                    await message.reply(f'{msg.jump_url}')
 
-            else:
+                LOGGER.info(f'Done')
-                await message.reply(f"NObody (in the past {days} days)...gah, leave me alone!")
-
-            LOGGER.info(f'Done')
 
     async def respond_to_joke(self, message: Message):
         for joke in self.jokes:
@@ -150,16 +151,6 @@ def get_emoji_name(string: str) -> str:
     return string.lower().strip()
 
 
-day_regex = re.compile('(?P<days>\d+) days')
-
-
 def get_days(input_str):
-    if (m := day_regex.search(input_str)):
+    if (m := re.search('(?P<days>\d+) days', input_str)):
         return int(m.group('days'))
-
-
-def filter_days(df: pd.DataFrame, days: int) -> pd.DataFrame:
-    start = (datetime.today() - timedelta(days=days)).astimezone()
-    valid_dates = df['datetime'] > start
-    df = df.loc[valid_dates]
-    return df

src/kwaylon/msg.py

@@ -9,13 +9,21 @@ from nextcord.utils import AsyncIterator
 LOGGER = logging.getLogger(__name__)
 
 
-async def message_gen(client: Client, limit: int = None, days: int = 90, **kwargs) -> AsyncIterator[Message]:
+async def message_gen(client: Client,
-    if 'after' not in kwargs:
+                      limit: int = None,
-        kwargs['after'] = (datetime.today() - timedelta(days=days))
+                      before: datetime = None,
-    elif isinstance((after := kwargs.get('after', None)), datetime):
+                      days: int = None,
-        kwargs['after'] = after.replace(tzinfo=None)
+                      after: datetime = None,
+                      oldest_first: bool = True) -> AsyncIterator[Message]:
+    if days is not None:
+        after = (datetime.today() - timedelta(days=days)).astimezone()
 
-    kwargs['limit'] = limit
+    kwargs = {
+        'limit': limit,
+        'before': before,
+        'after': after,
+        'oldest_first': oldest_first
+    }
 
     LOGGER.info(kwargs)
     for channel in client.get_all_channels():

src/kwaylon/reactions.py

@@ -1,6 +1,7 @@
 import logging
 import sqlite3
 from dataclasses import dataclass
+from datetime import datetime, timedelta
 from pathlib import Path
 
 import pandas as pd
@@ -19,18 +20,17 @@ class ReactionData:
         return sqlite3.connect(self.path, *args, **kwargs)
 
     async def scan_messages(self, client: Client, **kwargs):
-        try:
+        with self.connect() as con:
-            with self.connect() as con:
+            try:
                 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:
+            except Exception as e:
-            LOGGER.exception(e)
+                LOGGER.exception(e)
-        finally:
+        con.close()
-            con.close()
 
     def add_reactions_from_message(self, msg: Message, con: sqlite3.Connection = None):
-        con = con or sqlite3.connect(self.path)
+        con = con or self.connect()
 
         try:
             con.execute(f'DELETE FROM reactions WHERE msg_id = {msg.id}')
@@ -58,6 +58,25 @@ class ReactionData:
         if close:
             con.close()
 
-        res['datetime'] = pd.to_datetime(res['datetime'])
+        res['datetime'] = pd.to_datetime(res['datetime'], utc=True)
 
         return res.sort_values('count', ascending=False)
+
+    def row_count(self, con: sqlite3.Connection = None) -> int:
+        with con or self.connect() as con:
+            cur = con.execute('SELECT COUNT(*) FROM reactions')
+            n = cur.fetchone()[0]
+        con.close()
+        return n
+
+    def earliest(self, con: sqlite3.Connection = None):
+        with con or self.connect() as con:
+            cur = con.execute('SELECT MIN(datetime) FROM reactions')
+        date = pd.to_datetime(cur.fetchone()[0])
+        con.close()
+        return date
+
+
+def filter_days(df: pd.DataFrame, days: int) -> pd.DataFrame:
+    start = (datetime.today() - timedelta(days=days)).astimezone()
+    return df.loc[df['datetime'] > start]