From 1e1037dbe2f5297120330973b6bc8288ed8923ac Mon Sep 17 00:00:00 2001 From: jsl12 Date: Mon, 2 Aug 2021 15:03:13 -0500 Subject: [PATCH] broke out each load function --- msg.py | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/msg.py b/msg.py index 8f6e267..362103b 100644 --- a/msg.py +++ b/msg.py @@ -58,7 +58,7 @@ async def reaction_df(msgs: Iterable[discord.Message]): def add_reactions(con: sqlite3.Connection, new_reacts: pd.DataFrame, table_name: str = 'reactions'): - react_df = pd.read_sql(f'select * from {table_name}', con) + react_df = load_reactions(con, table_name) react_df = react_df.append(new_reacts, ignore_index=True) react_df = react_df.drop_duplicates(['msg id', 'emoji id']).reset_index(drop=True) try: @@ -71,7 +71,7 @@ def add_reactions(con: sqlite3.Connection, new_reacts: pd.DataFrame, table_name: def add_msgs(con: sqlite3.Connection, new_msgs: pd.DataFrame, table_name: str = 'msgs'): - msg_df = pd.read_sql(f'select * from {table_name}', con, index_col='id') + msg_df = load_msgs(con, table_name) msg_df = msg_df.append(new_msgs) msg_df['created'] = pd.to_datetime(msg_df['created'], utc=True) msg_df = msg_df[~msg_df.index.duplicated()].sort_values('created', ascending=False) @@ -103,11 +103,18 @@ async def get_and_save(db_file, client: discord.Client, limit: int, days: int): return msg_df, react_df -def load_both(con) -> Tuple[pd.DataFrame, pd.DataFrame]: - return ( - pd.read_sql(f'select * from msgs', con, index_col='id'), - pd.read_sql(f'select * from reactions', con) - ) +def load_both(con: sqlite3.Connection) -> Tuple[pd.DataFrame, pd.DataFrame]: + return (load_msgs(con), load_reactions(con)) + + +def load_msgs(con: sqlite3.Connection, table_name: str = 'msgs') -> pd.DataFrame: + df = pd.read_sql(f'select * from {table_name}', con, index_col='id') + df['created'] = df['created'].apply(pd.to_datetime, utc=True) + return df + + +def load_reactions(con: sqlite3.Connection, table_name: str = 'reactions') -> pd.DataFrame: + return pd.read_sql(f'select * from {table_name}', con) def cancellations(msg_df, react_df, days: int = 10) -> pd.DataFrame: @@ -143,7 +150,7 @@ def cancelled_totals(cdf: pd.DataFrame) -> pd.DataFrame: 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} total`\n{row['worst']}" + f"`{name.ljust(width + 1)}with {row['total']:<2} total`" for name, row in df.iterrows() )