improved reaction handling, added some comments

This commit is contained in:
2021-08-14 13:05:19 -05:00
parent 972b5afab5
commit fecbf5cb54
3 changed files with 77 additions and 71 deletions

64
data.py
View File

@@ -6,9 +6,8 @@ from pathlib import Path
import discord import discord
import pandas as pd 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__) LOGGER = logging.getLogger(__name__)
@@ -22,10 +21,9 @@ class MsgData:
async def create(cls, client: discord.Client, **kwargs): async def create(cls, client: discord.Client, **kwargs):
self = MsgData() self = MsgData()
self.lock = asyncio.Lock() self.lock = asyncio.Lock()
async with self.lock:
self.msgs: pd.DataFrame = await message_df(client, **kwargs) self.msgs: pd.DataFrame = await message_df(client, **kwargs)
self.msgs = self.msgs.sort_values('created') self.msgs = self.msgs.sort_values('created')
self.reactions: pd.DataFrame = await reaction_df(self.msgs['object'].tolist()) self.reactions: pd.DataFrame = full_reaction_df(self.msgs['object'].tolist())
return self return self
@classmethod @classmethod
@@ -56,7 +54,7 @@ class MsgData:
index=True, index=True,
index_label=self.msgs.index.name index_label=self.msgs.index.name
) )
self.reactions.to_sql( self.reactions.drop('object', axis=1).to_sql(
name='reactions', name='reactions',
con=con, con=con,
if_exists='replace', if_exists='replace',
@@ -83,34 +81,39 @@ class MsgData:
self.msgs.loc[message.id] = pd.Series(mdict) self.msgs.loc[message.id] = pd.Series(mdict)
LOGGER.info(f'Added message id {message.id} from {message.author}: {message.content}') LOGGER.info(f'Added message id {message.id} from {message.author}: {message.content}')
async def update_reaction(self, client: discord.Client, payload: RawReactionActionEvent): async def update_reaction(self, msg: discord.Message):
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 with self.lock: async with self.lock:
# Drop all the reactions for this message id, if there are any
try: 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: except KeyError as e:
LOGGER.warning(e) pass
if (new := await reaction_series(msg=msg)) is not None: # If there are reactions on the message after the change
new = new.set_index(['msg id', 'emoji']) if len(msg.reactions) > 0:
new = reaction_df(msg)
self.reactions = pd.concat([self.reactions, new]) 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: 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""" """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 # Get the ids of messages that that have the targeted emoji
count_ids = counts.index.drop_duplicates() message_id_counts: pd.Index = counts.index.drop_duplicates()
# filter by the messages that have actually been captured in the self.msgs DataFrame # There could be a situation where a message id in message_id_counts isn't actually in the self.msgs DataFrame
count_ids = count_ids[count_ids.isin(self.msgs.index.get_level_values(0))] # 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: # If there were actually some message ids found
res = self.msgs.loc[count_ids] if message_id_counts.shape[0] > 0:
res: pd.DataFrame = self.msgs.loc[message_id_counts]
res['count'] = counts res['count'] = counts
@@ -122,12 +125,13 @@ class MsgData:
raise KeyError(f'No messages found with {emoji_name} reactions') raise KeyError(f'No messages found with {emoji_name} reactions')
def emoji_counts(self, emoji_name: str) -> pd.Series: 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' assert isinstance(emoji_name, str), f'emoji_name must be a string'
try: 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: except KeyError as e:
LOGGER.error( LOGGER.error(f' {emoji_name} not found out of {self.unique_emojis.shape[0]} unique emojis')
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 raise
@property @property
@@ -135,7 +139,7 @@ class MsgData:
return self.reactions.index.get_level_values(1).drop_duplicates() return self.reactions.index.get_level_values(1).drop_duplicates()
def emoji_totals(self, emoji_name: str, days: int = None) -> pd.Series: 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 return (self
.emoji_messages(emoji_name, days) .emoji_messages(emoji_name, days)
.groupby('user id') .groupby('user id')
@@ -145,12 +149,10 @@ class MsgData:
async def emoji_leaderboard(self, client: discord.Client, emoji_name: str, days: int): async def emoji_leaderboard(self, client: discord.Client, emoji_name: str, days: int):
counts: pd.Series = self.emoji_totals(emoji_name, days) 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]) 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 = f'{emoji_name} totals, past {days} days\n'
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
def worst_offsenses(self, user: str, days: int): def worst_offsenses(self, user: str, days: int):

23
msg.py
View File

@@ -53,26 +53,21 @@ def message_dict(m: discord.Message) -> Dict:
} }
async def reaction_df(msgs: Iterable[discord.Message]): def full_reaction_df(msgs: Iterable[discord.Message]):
return pd.concat([await reaction_series(msg) for msg in msgs if len(msg.reactions) > 0]).set_index( return pd.concat([reaction_df(msg) for msg in msgs])
['msg id', 'emoji'])
async def reaction_series(msg: discord.Message): def reaction_df(msg: discord.Message):
if len(msg.reactions) > 0: df = pd.DataFrame([reaction_dict(r) for r in msg.reactions])
return pd.DataFrame([ return df.set_index(['msg id', 'emoji']) if not df.empty else df
await reaction_dict(r)
for r in msg.reactions
])
async def reaction_dict(r: discord.Reaction) -> Dict: def reaction_dict(r: discord.Reaction) -> Dict:
is_emoji = isinstance(r.emoji, (discord.Emoji, discord.PartialEmoji))
# LOGGER.info(repr(r.emoji))
return { return {
'object': r,
'msg id': r.message.id, 'msg id': r.message.id,
'emoji': r.emoji.name if is_emoji else r.emoji.encode('unicode-escape').decode('ascii'), 'emoji': r.emoji.name if r.is_custom_emoji() else r.emoji,
'emoji id': r.emoji.id if is_emoji else None, 'emoji id': r.emoji.id if r.is_custom_emoji() else None,
'count': int(r.count), 'count': int(r.count),
} }

View File

@@ -3,8 +3,9 @@ import logging
import os import os
import re import re
from pathlib import Path from pathlib import Path
from typing import Union
import discord import discord
from discord import RawReactionActionEvent, RawReactionClearEmojiEvent
from dotenv import load_dotenv from dotenv import load_dotenv
import data import data
@@ -47,9 +48,9 @@ class RoboPage(discord.Client):
self.data: data.MsgData = await data.MsgData.create( self.data: data.MsgData = await data.MsgData.create(
client=self, client=self,
limit=5000, # limit=5000,
# limit=20, limit=20,
days=30, # days=30,
) )
self.data.to_sql('messages.db') self.data.to_sql('messages.db')
LOGGER.info(f'{self.data.msgs.shape[0]} messages total') LOGGER.info(f'{self.data.msgs.shape[0]} messages total')
@@ -61,26 +62,23 @@ class RoboPage(discord.Client):
await self.data.add_msg(message) await self.data.add_msg(message)
if (m := self.leaderboard_regex.match(message.content)) is not None: if (m := self.leaderboard_regex.match(message.content)) is not None:
days = int(m.group('days')) or 14
emoji = m.group('emoji').lower()
try: try:
await message.reply(await self.data.emoji_leaderboard( await message.reply(
client=self, await self.data.emoji_leaderboard(client=self, emoji_name=emoji, days=days)
emoji_name=m.group('emoji').lower(), )
days=14
))
except KeyError as e: except KeyError as e:
LOGGER.exception(e) LOGGER.exception(e)
await message.reply(f"I couldn't find any {m.group('emoji')} reactions. Leave me alone!") await message.reply(f"I couldn't find any {m.group('emoji')} reactions. Leave me alone!")
return return
elif (m := self.most_regex.match(message.content)) is not None: 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: try:
await message.reply( await message.reply(
await self.data.biggest_single( await self.data.biggest_single(client=self, emoji=emoji, days=days))
client=self,
emoji=m.group('emoji').lower(),
days=int(days)
))
except IndexError as e: except IndexError as e:
await message.reply('NObody') await message.reply('NObody')
return return
@@ -90,6 +88,21 @@ class RoboPage(discord.Client):
LOGGER.info(f'{joke.__class__.__name__} detected: {message.content}, {scan_res.group()}') LOGGER.info(f'{joke.__class__.__name__} detected: {message.content}, {scan_res.group()}')
await joke.respond(message, self, scan_res) 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__': if __name__ == '__main__':
load_dotenv() load_dotenv()
@@ -110,17 +123,13 @@ if __name__ == '__main__':
@client.event @client.event
async def on_raw_reaction_add(payload): async def on_raw_reaction_add(payload: RawReactionActionEvent):
LOGGER.info(payload) await client.handle_raw_reaction(payload)
if hasattr(client, 'data'):
await client.data.update_reaction(payload=payload, client=client)
@client.event @client.event
async def on_raw_reaction_remove(payload): async def on_raw_reaction_remove(payload: RawReactionClearEmojiEvent):
LOGGER.info(payload) await client.handle_raw_reaction(payload)
if hasattr(client, 'data'):
await client.data.update_reaction(payload=payload, client=client)
client.run() client.run()