diff --git a/data.py b/data.py new file mode 100644 index 0000000..9d746d7 --- /dev/null +++ b/data.py @@ -0,0 +1,144 @@ +import asyncio +import logging +import sqlite3 +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, emoji_messages, \ + emoji_totals + +LOGGER = logging.getLogger(__name__) + + +class MsgData: + msgs: pd.DataFrame + reactions: pd.DataFrame + lock: asyncio.Lock + + @classmethod + async def create(cls, client: discord.Client, **kwargs): + self = MsgData() + 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()) + self.lock = asyncio.Lock() + return self + + @classmethod + def from_sql(cls, db): + if isinstance(db, (str, Path)): + con = sqlite3.connect(db) + elif isinstance(db, sqlite3.Connection): + con = db + + self = MsgData() + self.msgs: pd.DataFrame = pd.read_sql('select * from msgs', con=con, index_col='id') + self.msgs['created'] = self.msgs['created'].apply(pd.to_datetime, utc=True) + self.reactions: pd.DataFrame = pd.read_sql('select * from reactions', con).set_index(['msg id', 'emoji']) + return self + + def to_sql(self, db): + if isinstance(db, (str, Path)): + con = sqlite3.connect(db) + elif isinstance(db, sqlite3.Connection): + con = db + + self.msgs.drop('object', axis=1).to_sql( + name='msgs', + con=con, + if_exists='replace', + index=True, + index_label=self.msgs.index.name + ) + self.reactions.to_sql( + name='reactions', + con=con, + if_exists='replace', + index=True, + index_label=self.reactions.index.name + ) + + def __str__(self): + return str(self.msgs) + '\n\n' + str(self.reactions) + + async def add_msg(self, message: discord.Message): + async with self.lock: + mdict = message_dict(message) + mdict.pop('id') + 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 with self.lock: + try: + self.reactions.drop(msg.id, level=0, axis=0) + except KeyError as e: + LOGGER.warning(e) + + if (new := await reaction_series(msg=msg)) is not None: + self.reactions = pd.concat([self.reactions, new.set_index(['msg id', 'emoji'])]) + LOGGER.info(f'\n{str(new)}') + + def emoji_messages(self, emoji_name: str, days: int): + res = emoji_messages(msg_df=self.msgs, react_df=self.reactions, emoji_name=emoji_name, days=days) + if res is None: + raise KeyError(f'No emojis found for {emoji_name}') + else: + return res + + def emoji_totals(self, emoji_name: str, days: int): + return emoji_totals(edf=self.emoji_messages(emoji_name, days)) + + def emoji_leaderboard(self, emoji_name: str, days: int): + df = self.emoji_totals(emoji_name, days) + width = max(list(map(lambda s: len(str(s)), df.index.values))) + res = f'{emoji_name} totals, past {days} days\n' + res += '\n'.join( + f"`{str(name).ljust(width + 1)}with {row['total']:<2.0f} total`" + for name, row in df.iterrows() + ) + return res + + def cancellation_leaderboard(self, days): + return self.emoji_leaderboard(emoji_name='cancelled', days=days) + + def worst_offsenses(self, user: str, days: int): + cdf = self.emoji_messages('cancelled', days=days) + cdf = cdf[cdf['display_name'].str.contains(user, case=False)] + + if cdf.shape[0] > 0: + res = f'{user}\'s top 5 cancellations in the last {days} days:\n' + res += f'\n'.join( + f'`{row["count"]:<2.0f}cancellations`\n{row["link"]}' for idx, row in cdf.iloc[:5].iterrows()) + else: + res = f'No cancellations for {user} in the past {days} days' + + return res + + def biggest_daddy(self, days, top: int = None): + df = self.emoji_totals('daddy', days) + if top is not None: + df = df.iloc[:top] + width = max(list(map(lambda s: len(str(s)), df.index.values))) + res = f'Daddy totals, past {days} days\n' + res += '\n'.join( + f"`{name.ljust(width + 1)}with {row['total']:<2.0f} total`" + for name, row in df.iterrows() + ) + return res + + async def biggest_single(self, client: discord.Client, emoji: str, days: int) -> str: + data = self.emoji_totals(emoji_name=emoji, days=days) + username = data.index[0] + reacted_msgs = self.emoji_messages(emoji_name=emoji, days=days) + d = reacted_msgs.set_index('display_name')['user id'].drop_duplicates().to_dict() + user: discord.User = await client.fetch_user(user_id=d[username]) + LOGGER.info(f'User: {user.mention}') + return f'{user.mention} with {data.iloc[0]["total"]:.0f} over the past {int(days)} days' diff --git a/msg.py b/msg.py index 349c37e..abcb5a3 100644 --- a/msg.py +++ b/msg.py @@ -1,150 +1,15 @@ -import asyncio import logging import os -import sqlite3 from datetime import datetime, timedelta -from pathlib import Path from typing import Dict, Iterable import discord import pandas as pd -from discord.raw_models import RawReactionActionEvent from dotenv import load_dotenv LOGGER = logging.getLogger(__name__) -class MsgData: - msgs: pd.DataFrame - reactions: pd.DataFrame - lock: asyncio.Lock - - @classmethod - async def create(cls, client: discord.Client, **kwargs): - self = MsgData() - 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()) - self.lock = asyncio.Lock() - return self - - @classmethod - def from_sql(cls, db): - if isinstance(db, (str, Path)): - con = sqlite3.connect(db) - elif isinstance(db, sqlite3.Connection): - con = db - - self = MsgData() - self.msgs: pd.DataFrame = pd.read_sql('select * from msgs', con=con, index_col='id') - self.msgs['created'] = self.msgs['created'].apply(pd.to_datetime, utc=True) - self.reactions: pd.DataFrame = pd.read_sql('select * from reactions', con).set_index(['msg id', 'emoji']) - return self - - def to_sql(self, db): - if isinstance(db, (str, Path)): - con = sqlite3.connect(db) - elif isinstance(db, sqlite3.Connection): - con = db - - self.msgs.drop('object', axis=1).to_sql( - name='msgs', - con=con, - if_exists='replace', - index=True, - index_label=self.msgs.index.name - ) - self.reactions.to_sql( - name='reactions', - con=con, - if_exists='replace', - index=True, - index_label=self.reactions.index.name - ) - - def __str__(self): - return str(self.msgs) + '\n\n' + str(self.reactions) - - async def add_msg(self, message: discord.Message): - async with self.lock: - mdict = message_dict(message) - mdict.pop('id') - 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) - - with self.lock: - try: - self.reactions.drop(msg.id, level=0, axis=0) - except KeyError as e: - LOGGER.warning(e) - - if (new := await reaction_series(msg=msg)) is not None: - self.reactions = pd.concat([self.reactions, new.set_index(['msg id', 'emoji'])]) - LOGGER.info(f'\n{str(new)}') - - def emoji_messages(self, emoji_name: str, days: int): - res = emoji_messages(msg_df=self.msgs, react_df=self.reactions, emoji_name=emoji_name, days=days) - if res is None: - raise KeyError(f'No emojis found for {emoji_name}') - else: - return res - - def emoji_totals(self, emoji_name: str, days: int): - return emoji_totals(edf=self.emoji_messages(emoji_name, days)) - - def emoji_leaderboard(self, emoji_name: str, days: int): - df = self.emoji_totals(emoji_name, days) - width = max(list(map(lambda s: len(str(s)), df.index.values))) - res = f'{emoji_name} totals, past {days} days\n' - res += '\n'.join( - f"`{str(name).ljust(width + 1)}with {row['total']:<2.0f} total`" - for name, row in df.iterrows() - ) - return res - - def cancellation_leaderboard(self, days): - return self.emoji_leaderboard(emoji_name='cancelled', days=days) - - def worst_offsenses(self, user: str, days: int): - cdf = self.emoji_messages('cancelled', days=days) - cdf = cdf[cdf['display_name'].str.contains(user, case=False)] - - if cdf.shape[0] > 0: - res = f'{user}\'s top 5 cancellations in the last {days} days:\n' - res += f'\n'.join( - f'`{row["count"]:<2.0f}cancellations`\n{row["link"]}' for idx, row in cdf.iloc[:5].iterrows()) - else: - res = f'No cancellations for {user} in the past {days} days' - - return res - - def biggest_daddy(self, days, top: int = None): - df = self.emoji_totals('daddy', days) - if top is not None: - df = df.iloc[:top] - width = max(list(map(lambda s: len(str(s)), df.index.values))) - res = f'Daddy totals, past {days} days\n' - res += '\n'.join( - f"`{name.ljust(width + 1)}with {row['total']:<2.0f} total`" - for name, row in df.iterrows() - ) - return res - - async def biggest_single(self, client: discord.Client, emoji: str, days: int) -> str: - data = self.emoji_totals(emoji_name=emoji, days=days) - username = data.index[0] - reacted_msgs = self.emoji_messages(emoji_name=emoji, days=days) - d = reacted_msgs.set_index('display_name')['user id'].drop_duplicates().to_dict() - user: discord.User = await client.fetch_user(user_id=d[username]) - LOGGER.info(f'User: {user.mention}') - return f'{user.mention} with {data.iloc[0]["total"]:.0f} over the past {int(days)} days' - - def convert_emoji(emoji): try: emoji.name.encode('ascii') diff --git a/robopage.py b/robopage.py index 5e9bdb8..89ccab2 100644 --- a/robopage.py +++ b/robopage.py @@ -6,8 +6,8 @@ from threading import Lock import discord from dotenv import load_dotenv +import data import jokes -import msg logging.basicConfig(level=logging.INFO) @@ -37,7 +37,7 @@ class RoboPage(discord.Client): channel: discord.TextChannel = discord.utils.get(self.get_all_channels(), name='robotics-facility') await channel.send(f"I'm aliiiiiive {discord.utils.get(self.emojis, name='kaylon')}") - self.data: msg.MsgData = await msg.MsgData.create( + self.data: data.MsgData = await data.MsgData.create( client=self, limit=3000, # limit=20,