added from_sql and to_sql methods to MsgData

This commit is contained in:
2021-08-10 22:08:15 -05:00
parent 156715879f
commit 6a359d85ff
2 changed files with 62 additions and 31 deletions

71
msg.py
View File

@@ -1,7 +1,9 @@
import asyncio import asyncio
import logging import logging
import os import os
import sqlite3
from datetime import datetime, timedelta from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, Iterable from typing import Dict, Iterable
import discord import discord
@@ -26,6 +28,40 @@ class MsgData:
self.lock = asyncio.Lock() self.lock = asyncio.Lock()
return self 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): def __str__(self):
return str(self.msgs) + '\n\n' + str(self.reactions) return str(self.msgs) + '\n\n' + str(self.reactions)
@@ -63,9 +99,21 @@ class MsgData:
def cancellations(self, days: int = 14): def cancellations(self, days: int = 14):
return cancellations(msg_df=self.msgs, react_df=self.reactions, days=days) 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)) 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): async def message_df(client: discord.Client, **kwargs):
return pd.DataFrame( return pd.DataFrame(
@@ -94,6 +142,7 @@ def message_dict(m: discord.Message) -> Dict:
'id': m.id, 'id': m.id,
'created': m.created_at.astimezone(), 'created': m.created_at.astimezone(),
'display_name': m.author.display_name, 'display_name': m.author.display_name,
'user id': m.author.id,
'message': m.content, 'message': m.content,
'channel': m.channel.name, 'channel': m.channel.name,
'channel link': f'<#{m.channel.id}>', '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): async def reaction_series(msg: discord.Message):
return pd.DataFrame( return pd.DataFrame([
[ await reaction_dict(r)
await reaction_dict(r) for r in msg.reactions
for r in msg.reactions if isinstance(r.emoji, discord.Emoji)
if isinstance(r.emoji, discord.Emoji) ])
]
)
async def reaction_dict(r: discord.Reaction) -> Dict: 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) }).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__': if __name__ == '__main__':
client = discord.Client() client = discord.Client()

View File

@@ -6,8 +6,8 @@ from threading import Lock
import discord import discord
from dotenv import load_dotenv from dotenv import load_dotenv
import jokes
import msg import msg
from jokes import CumJoke, BlackJoke, AssJoke, DominosJoke
logging.basicConfig(level=logging.INFO) logging.basicConfig(level=logging.INFO)
@@ -21,12 +21,9 @@ class RoboPage(discord.Client):
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
super(RoboPage, self).__init__(*args, **kwargs) super(RoboPage, self).__init__(*args, **kwargs)
self.jokes = [ attrs = filter(lambda n: n.endswith('Joke') and not n.startswith('Joke'), dir(jokes))
CumJoke(), attrs = map(lambda n: getattr(jokes, n)(), attrs)
BlackJoke(), self.jokes = list(attrs)
AssJoke(),
DominosJoke()
]
self.lock = Lock() self.lock = Lock()
def run(self): def run(self):
@@ -40,9 +37,9 @@ class RoboPage(discord.Client):
self.data: msg.MsgData = await msg.MsgData.create( self.data: msg.MsgData = await msg.MsgData.create(
client=self, client=self,
limit=3000, limit=3000,
# limit=20,
days=14, days=14,
) )
self.data.to_sql('messages.db')
LOGGER.info(str(self.data.msgs.columns)) LOGGER.info(str(self.data.msgs.columns))
LOGGER.info(str(self.data.reactions.columns)) LOGGER.info(str(self.data.reactions.columns))
@@ -51,7 +48,7 @@ class RoboPage(discord.Client):
if message.author != self.user: if message.author != self.user:
if 'most cancelled' in message.content: 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<name>\w+)', message.content)) is not None: elif (m := re.search('top cancelled (?P<name>\w+)', message.content)) is not None:
async with self.data.lock: 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}') print(f'{joke.__class__.__name__} detected:\n{message.content}\n{scan_res}')
await joke.respond(message, self, 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): def top_cancellations(self, user: str, days: int):
cdf = self.data.cancellations(days) cdf = self.data.cancellations(days)
cdf = cdf[cdf['display_name'].str.contains(user, case=False)] cdf = cdf[cdf['display_name'].str.contains(user, case=False)]
@@ -88,7 +80,7 @@ if __name__ == '__main__':
@client.event @client.event
async def on_ready(): async def on_ready():
print(f'{client.user} has connected to Discord!') print(len(list(client.get_all_members())))
await client.handle_ready() await client.handle_ready()
print(client.data.cancellation_totals(14)) print(client.data.cancellation_totals(14))