diff --git a/data.py b/data.py index 738f6f5..43cd372 100644 --- a/data.py +++ b/data.py @@ -6,9 +6,8 @@ from pathlib import Path import discord import pandas as pd -from discord import RawReactionActionEvent -from msg import message_df, reaction_df, message_dict, LOGGER, convert_emoji, reaction_series +from msg import message_df, full_reaction_df, message_dict, LOGGER, reaction_df LOGGER = logging.getLogger(__name__) @@ -22,11 +21,10 @@ class MsgData: async def create(cls, client: discord.Client, **kwargs): self = MsgData() self.lock = asyncio.Lock() - async with self.lock: - self.msgs: pd.DataFrame = await message_df(client, **kwargs) - self.msgs = self.msgs.sort_values('created') - self.reactions: pd.DataFrame = await reaction_df(self.msgs['object'].tolist()) - return self + self.msgs: pd.DataFrame = await message_df(client, **kwargs) + self.msgs = self.msgs.sort_values('created') + self.reactions: pd.DataFrame = full_reaction_df(self.msgs['object'].tolist()) + return self @classmethod def from_sql(cls, db, local_tz='US/Central'): @@ -56,7 +54,7 @@ class MsgData: index=True, index_label=self.msgs.index.name ) - self.reactions.to_sql( + self.reactions.drop('object', axis=1).to_sql( name='reactions', con=con, if_exists='replace', @@ -83,34 +81,39 @@ class MsgData: self.msgs.loc[message.id] = pd.Series(mdict) LOGGER.info(f'Added message id {message.id} from {message.author}: {message.content}') - async def update_reaction(self, client: discord.Client, payload: RawReactionActionEvent): - payload.emoji: discord.PartialEmoji = convert_emoji(payload.emoji) - chan: discord.TextChannel = await client.fetch_channel(channel_id=payload.channel_id) - msg: discord.Message = await chan.fetch_message(payload.message_id) - + async def update_reaction(self, msg: discord.Message): async with self.lock: + # Drop all the reactions for this message id, if there are any try: - self.reactions.drop(msg.id, level=0, axis=0) + self.reactions.drop(msg.id, level=0, axis=0, inplace=True) except KeyError as e: - LOGGER.warning(e) + pass - if (new := await reaction_series(msg=msg)) is not None: - new = new.set_index(['msg id', 'emoji']) + # If there are reactions on the message after the change + if len(msg.reactions) > 0: + new = reaction_df(msg) self.reactions = pd.concat([self.reactions, new]) - LOGGER.info(f'\n{str(new)}') + LOGGER.info(str(new.droplevel(level=0, axis=0).loc[:, 'count'])) + + if msg.id not in self.msgs.index: + await self.add_msg(msg) + + return new def emoji_messages(self, emoji_name: str, days: int = None) -> pd.DataFrame: """Creates a DataFrame of the messages that have reactions with a certain emoji. Includes a 'count' column""" - counts = self.emoji_counts(emoji_name) + counts: pd.DataFrame = self.emoji_counts(emoji_name) - # get the ids of messages that that have the targeted emoji - count_ids = counts.index.drop_duplicates() + # Get the ids of messages that that have the targeted emoji + message_id_counts: pd.Index = counts.index.drop_duplicates() - # filter by the messages that have actually been captured in the self.msgs DataFrame - count_ids = count_ids[count_ids.isin(self.msgs.index.get_level_values(0))] + # There could be a situation where a message id in message_id_counts isn't actually in the self.msgs DataFrame + # Filter to keep only the messages that have actually been captured in the self.msgs DataFrame + message_id_counts: pd.Index = message_id_counts[message_id_counts.isin(self.msgs.index.get_level_values(0))] - if count_ids.shape[0] > 0: - res = self.msgs.loc[count_ids] + # If there were actually some message ids found + if message_id_counts.shape[0] > 0: + res: pd.DataFrame = self.msgs.loc[message_id_counts] res['count'] = counts @@ -122,12 +125,13 @@ class MsgData: raise KeyError(f'No messages found with {emoji_name} reactions') def emoji_counts(self, emoji_name: str) -> pd.Series: + """Creates a Series indexed by message id and with the number of reactions with emoji_name as values""" assert isinstance(emoji_name, str), f'emoji_name must be a string' try: - return self.reactions.loc[pd.IndexSlice[:, emoji_name],'count'].droplevel(1).sort_values(ascending=False) + return self.reactions.loc[pd.IndexSlice[:, emoji_name], 'count'].droplevel(1).sort_values(ascending=False) except KeyError as e: - LOGGER.error( - f' {emoji_name} not found out of {self.unique_emojis.shape[0]} unique emojis') + LOGGER.error(f' {emoji_name} not found out of {self.unique_emojis.shape[0]} unique emojis') + LOGGER.error(f'{self.reactions.index.get_level_values(1)}') raise @property @@ -135,7 +139,7 @@ class MsgData: return self.reactions.index.get_level_values(1).drop_duplicates() def emoji_totals(self, emoji_name: str, days: int = None) -> pd.Series: - """Creates a Series of the counts for each user id""" + """Creates a Series indexed by user id and with the number of reactions with emoji_name as values""" return (self .emoji_messages(emoji_name, days) .groupby('user id') @@ -145,12 +149,10 @@ class MsgData: async def emoji_leaderboard(self, client: discord.Client, emoji_name: str, days: int): counts: pd.Series = self.emoji_totals(emoji_name, days) counts.index = pd.Index([(await client.fetch_user(user_id=uid)).display_name for uid in counts.index]) - width = max(list(map(lambda s: len(str(s)), counts.index.values))) + width = max([len(str(s)) for s in counts.index.values]) res = f'{emoji_name} totals, past {days} days\n' - res += '\n'.join( - f"`{str(name).ljust(width + 1)}with {cnt:<2.0f} total`" - for name, cnt in counts.iteritems() - ) + res += '\n'.join(f"`{str(name).ljust(width + 1)}with {cnt:<2.0f} total`" + for name, cnt in counts.iteritems()) return res def worst_offsenses(self, user: str, days: int): diff --git a/msg.py b/msg.py index 7de40f0..9562eb4 100644 --- a/msg.py +++ b/msg.py @@ -53,26 +53,21 @@ def message_dict(m: discord.Message) -> Dict: } -async def reaction_df(msgs: Iterable[discord.Message]): - return pd.concat([await reaction_series(msg) for msg in msgs if len(msg.reactions) > 0]).set_index( - ['msg id', 'emoji']) +def full_reaction_df(msgs: Iterable[discord.Message]): + return pd.concat([reaction_df(msg) for msg in msgs]) -async def reaction_series(msg: discord.Message): - if len(msg.reactions) > 0: - return pd.DataFrame([ - await reaction_dict(r) - for r in msg.reactions - ]) +def reaction_df(msg: discord.Message): + df = pd.DataFrame([reaction_dict(r) for r in msg.reactions]) + return df.set_index(['msg id', 'emoji']) if not df.empty else df -async def reaction_dict(r: discord.Reaction) -> Dict: - is_emoji = isinstance(r.emoji, (discord.Emoji, discord.PartialEmoji)) - # LOGGER.info(repr(r.emoji)) +def reaction_dict(r: discord.Reaction) -> Dict: return { + 'object': r, 'msg id': r.message.id, - 'emoji': r.emoji.name if is_emoji else r.emoji.encode('unicode-escape').decode('ascii'), - 'emoji id': r.emoji.id if is_emoji else None, + 'emoji': r.emoji.name if r.is_custom_emoji() else r.emoji, + 'emoji id': r.emoji.id if r.is_custom_emoji() else None, 'count': int(r.count), } diff --git a/robopage.py b/robopage.py index b4b790b..b6eccf2 100644 --- a/robopage.py +++ b/robopage.py @@ -3,8 +3,9 @@ import logging import os import re from pathlib import Path - +from typing import Union import discord +from discord import RawReactionActionEvent, RawReactionClearEmojiEvent from dotenv import load_dotenv import data @@ -47,9 +48,9 @@ class RoboPage(discord.Client): self.data: data.MsgData = await data.MsgData.create( client=self, - limit=5000, - # limit=20, - days=30, + # limit=5000, + limit=20, + # days=30, ) self.data.to_sql('messages.db') LOGGER.info(f'{self.data.msgs.shape[0]} messages total') @@ -61,26 +62,23 @@ class RoboPage(discord.Client): await self.data.add_msg(message) if (m := self.leaderboard_regex.match(message.content)) is not None: + days = int(m.group('days')) or 14 + emoji = m.group('emoji').lower() try: - await message.reply(await self.data.emoji_leaderboard( - client=self, - emoji_name=m.group('emoji').lower(), - days=14 - )) + await message.reply( + await self.data.emoji_leaderboard(client=self, emoji_name=emoji, days=days) + ) except KeyError as e: LOGGER.exception(e) await message.reply(f"I couldn't find any {m.group('emoji')} reactions. Leave me alone!") return elif (m := self.most_regex.match(message.content)) is not None: - days = m.group('days') or 14 + days = int(m.group('days')) or 14 + emoji = m.group('emoji').lower() try: await message.reply( - await self.data.biggest_single( - client=self, - emoji=m.group('emoji').lower(), - days=int(days) - )) + await self.data.biggest_single(client=self, emoji=emoji, days=days)) except IndexError as e: await message.reply('NObody') return @@ -90,6 +88,21 @@ class RoboPage(discord.Client): LOGGER.info(f'{joke.__class__.__name__} detected: {message.content}, {scan_res.group()}') await joke.respond(message, self, scan_res) + async def handle_raw_reaction(self, payload: Union[RawReactionActionEvent, RawReactionClearEmojiEvent]): + LOGGER.info(payload) + guild = await client.fetch_guild(payload.guild_id) + channel = await guild.fetch_channel(payload.channel_id) + message = await channel.fetch_message(payload.message_id) + + if isinstance(payload, RawReactionActionEvent): + LOGGER.info( + f'{payload.member.display_name} added {payload.emoji} to\n{message.author.display_name}: {message.content}') + elif isinstance(payload, RawReactionClearEmojiEvent): + LOGGER.info(f'{payload.emoji} removed from\n{message.author.display_name}: {message.content}') + + if hasattr(self, 'data'): + await self.data.update_reaction(msg=message) + if __name__ == '__main__': load_dotenv() @@ -110,17 +123,13 @@ if __name__ == '__main__': @client.event - async def on_raw_reaction_add(payload): - LOGGER.info(payload) - if hasattr(client, 'data'): - await client.data.update_reaction(payload=payload, client=client) + async def on_raw_reaction_add(payload: RawReactionActionEvent): + await client.handle_raw_reaction(payload) @client.event - async def on_raw_reaction_remove(payload): - LOGGER.info(payload) - if hasattr(client, 'data'): - await client.data.update_reaction(payload=payload, client=client) + async def on_raw_reaction_remove(payload: RawReactionClearEmojiEvent): + await client.handle_raw_reaction(payload) client.run()