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, *args, **kwargs):
+    def __init__(self, limit: int = 5000, days: int = 30, db_path: Path = None, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        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,6 +80,7 @@ class Kwaylon(Client):
             emoji_name = get_emoji_name(emoji_ref)
             LOGGER.info(f'Most {emoji_name}')
 
+            async with message.channel.typing():
                 with self.data.connect() as con:
                     df = self.data.read_emoji(emoji_name, con)
                 con.close()
@@ -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]:
-    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 def message_gen(client: Client,
+                      limit: int = None,
+                      before: datetime = None,
+                      days: int = 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:
+            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:
                 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)
+        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]