From 6a359d85ffa56f2a5d8281f49721c71526ce427a Mon Sep 17 00:00:00 2001 From: jsl12 Date: Tue, 10 Aug 2021 22:08:15 -0500 Subject: [PATCH] added from_sql and to_sql methods to MsgData --- msg.py | 71 +++++++++++++++++++++++++++++++++++++++++------------ robopage.py | 22 ++++++----------- 2 files changed, 62 insertions(+), 31 deletions(-) diff --git a/msg.py b/msg.py index 3a6691a..f8d4220 100644 --- a/msg.py +++ b/msg.py @@ -1,7 +1,9 @@ import asyncio import logging import os +import sqlite3 from datetime import datetime, timedelta +from pathlib import Path from typing import Dict, Iterable import discord @@ -26,6 +28,40 @@ class MsgData: 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) @@ -63,9 +99,21 @@ class MsgData: def cancellations(self, days: int = 14): return cancellations(msg_df=self.msgs, react_df=self.reactions, days=days) - def cancellation_totals(self, days: int = 14): + def cancellation_totals(self, days): return cancelled_totals(cdf=self.cancellations(days=days)) + def cancellation_leaderboard(self, days, top: int = None): + df = self.cancellation_totals(days) + if top is not None: + df = df.iloc[:top] + width = max(list(map(lambda s: len(str(s)), df.index.values))) + res = f'Cancellation 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 message_df(client: discord.Client, **kwargs): return pd.DataFrame( @@ -94,6 +142,7 @@ def message_dict(m: discord.Message) -> Dict: 'id': m.id, 'created': m.created_at.astimezone(), 'display_name': m.author.display_name, + 'user id': m.author.id, 'message': m.content, 'channel': m.channel.name, 'channel link': f'<#{m.channel.id}>', @@ -106,13 +155,11 @@ async def reaction_df(msgs: Iterable[discord.Message]): async def reaction_series(msg: discord.Message): - return pd.DataFrame( - [ - await reaction_dict(r) - for r in msg.reactions - if isinstance(r.emoji, discord.Emoji) - ] - ) + return pd.DataFrame([ + await reaction_dict(r) + for r in msg.reactions + if isinstance(r.emoji, discord.Emoji) + ]) async def reaction_dict(r: discord.Reaction) -> Dict: @@ -160,14 +207,6 @@ def cancelled_totals(cdf: pd.DataFrame) -> pd.DataFrame: }).sort_values('total', ascending=False) -def report_string(df): - width = max(list(map(lambda s: len(str(s)), df.index.values))) - return '\n'.join( - f"`{name.ljust(width + 1)}with {row['total']:<2.0f} total`" - for name, row in df.iterrows() - ) - - if __name__ == '__main__': client = discord.Client() diff --git a/robopage.py b/robopage.py index 5b447a5..91b75af 100644 --- a/robopage.py +++ b/robopage.py @@ -6,8 +6,8 @@ from threading import Lock import discord from dotenv import load_dotenv +import jokes import msg -from jokes import CumJoke, BlackJoke, AssJoke, DominosJoke logging.basicConfig(level=logging.INFO) @@ -21,12 +21,9 @@ class RoboPage(discord.Client): def __init__(self, *args, **kwargs): super(RoboPage, self).__init__(*args, **kwargs) - self.jokes = [ - CumJoke(), - BlackJoke(), - AssJoke(), - DominosJoke() - ] + attrs = filter(lambda n: n.endswith('Joke') and not n.startswith('Joke'), dir(jokes)) + attrs = map(lambda n: getattr(jokes, n)(), attrs) + self.jokes = list(attrs) self.lock = Lock() def run(self): @@ -40,9 +37,9 @@ class RoboPage(discord.Client): self.data: msg.MsgData = await msg.MsgData.create( client=self, limit=3000, - # limit=20, days=14, ) + self.data.to_sql('messages.db') LOGGER.info(str(self.data.msgs.columns)) LOGGER.info(str(self.data.reactions.columns)) @@ -51,7 +48,7 @@ class RoboPage(discord.Client): if message.author != self.user: if 'most cancelled' in message.content: - await message.reply(self.get_cancelled_totals(days=14)) + await message.reply(self.data.cancellation_leaderboard(14)) elif (m := re.search('top cancelled (?P\w+)', message.content)) is not None: async with self.data.lock: @@ -62,11 +59,6 @@ class RoboPage(discord.Client): print(f'{joke.__class__.__name__} detected:\n{message.content}\n{scan_res}') await joke.respond(message, self, scan_res) - def get_cancelled_totals(self, days): - res = self.data.cancellation_totals(days) - res = f'Cancellation totals, past {days} days\n' + msg.report_string(res.iloc[:5]) - return res - def top_cancellations(self, user: str, days: int): cdf = self.data.cancellations(days) cdf = cdf[cdf['display_name'].str.contains(user, case=False)] @@ -88,7 +80,7 @@ if __name__ == '__main__': @client.event async def on_ready(): - print(f'{client.user} has connected to Discord!') + print(len(list(client.get_all_members()))) await client.handle_ready() print(client.data.cancellation_totals(14))