diff --git a/kwaylon/kwaylon.py b/kwaylon/kwaylon.py index 405eb15..aeb7691 100644 --- a/kwaylon/kwaylon.py +++ b/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\S+)') self.leaderboard_regex = re.compile( @@ -38,11 +40,9 @@ class Kwaylon(Client): # await alive() - self.data = MsgData(self.db_path) - - await self.data.load_sql() - if not hasattr(self.data, 'reactions'): - await self.data.scan_messages(client=self, limit=self.limit, days=self.days) + self.data = ReactionData(self.db_path) + # await self.data.scan_messages(client=self, limit=100) + self.data.read_all() 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\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) - await message.reply(res) + if (board := await self.leaderboard(df, **kwargs)) is not None: + 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() - counts = df['count'].sort_values(ascending=False) - counts.index = [(await self.fetch_user(idx)).display_name for idx in counts.index] + if df.shape[0] > 0: + df = df.groupby('auth_id').sum() + 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( - f"`{str(name).ljust(width + 1)}with {cnt:<2.0f} total`" - for name, cnt in counts.iteritems() - ) - return res + res = '\n'.join( + f"`{str(name).ljust(width + 1)}with {cnt:<2.0f} total`" + for name, cnt in counts.iteritems() + ) + 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': - LOGGER.info( - f'{payload.member.display_name} added {payload.emoji} to\n' + \ - f'{message.author.display_name}: {message.content}') - # 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) + with self.lock: + with self.data.connect() as con: + self.data.add_reactions_from_message(message, con) + con.close() async def fetch_message(self, row: pd.Series): guild = await self.fetch_guild(row['guild_id']) diff --git a/kwaylon/reactions.py b/kwaylon/reactions.py new file mode 100644 index 0000000..ee8474a --- /dev/null +++ b/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) \ No newline at end of file