diff --git a/data.py b/data.py index a920293..27b80e0 100644 --- a/data.py +++ b/data.py @@ -64,6 +64,15 @@ class MsgData: def __str__(self): return str(self.msgs) + '\n\n' + str(self.reactions) + def __repr__(self): + return f'<{__name__}.{self.__class__.__name__} with {self.msgs.shape[0]} messages and {self.reactions.shape[0]} reactions>' + + def __getitem__(self, item): + if isinstance(item, str): + return self.reactions.loc[pd.IndexSlice[:, item],].fillna(0).applymap(int) + elif isinstance(item, int): + return self.reactions.loc[pd.IndexSlice[item, :],].fillna(0).applymap(int) + async def add_msg(self, message: discord.Message): async with self.lock: mdict = message_dict(message)