From 014a792ad4f94f87dfbbbb6ae8ad4e46de9b088e Mon Sep 17 00:00:00 2001 From: klemek Date: Wed, 6 Jan 2021 16:30:53 +0100 Subject: [PATCH] log management to improve speed --- .gitignore | 1 + emotes.py | 2 +- log_manager.py | 142 +++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 log_manager.py diff --git a/.gitignore b/.gitignore index 738e240..bb73726 100755 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ __pycache__ .env error_* *.log +logs/ diff --git a/emotes.py b/emotes.py index 9a595ac..8b06ce0 100644 --- a/emotes.py +++ b/emotes.py @@ -73,7 +73,7 @@ EMOJI_REGEX = re.compile("(|:\\w+:)") def load_emojis(): global GLOBAL_EMOJIS, INV_GLOBAL_EMOJIS, EMOJI_REGEX emoji_list = [] - with open("emoji.json") as f: + with open("emoji.json", mode="r") as f: emoji_list = json.loads(f.readline().strip()) for emoji in EXTRA_EMOJI: emoji_list += [{"short_name": emoji, "unified": EXTRA_EMOJI[emoji]}] diff --git a/log_manager.py b/log_manager.py new file mode 100644 index 0000000..c089973 --- /dev/null +++ b/log_manager.py @@ -0,0 +1,142 @@ +from typing import Union, List +import os +import discord +import json + +LOG_DIR = "logs" + +if not os.path.exists(LOG_DIR): + os.mkdir(LOG_DIR) + + +CHUNK_SIZE = 1000 + + +class MessageLog: + def __init__(self, message: Union[discord.Message, dict]): + if isinstance(message, discord.Message): + self.id = message.id + self.created_at = message.created_at + self.edited_at = message.edited_at + self.author = message.author + self.pinned = message.pinned + self.mention_everyone = message.mention_everyone + self.tts = message.tts + self.reference = message.reference.id + self.content = message.content + self.mentions = message.raw_mentions + self.role_mentions = message.raw_role_mentions + self.channel_mentions = message.raw_channel_mentions + self.reactions = {} + elif isinstance(message, dict): + self.id = message["id"] + self.created_at = message["created_at"] + self.edited_at = message["edited_at"] + self.author = message["author"] + self.pinned = message["pinned"] + self.mention_everyone = message["mention_everyone"] + self.tts = message["tts"] + self.reference = message["reference.id"] + self.content = message["content"] + self.mentions = message["raw_mentions"] + self.role_mentions = message["raw_role_mentions"] + self.channel_mentions = message["raw_channel_mentions"] + self.reactions = message["reactions"] + + async def load(self, message: discord.Message): + for reaction in message.reactions: + self.reactions[str(reaction)] = [] + async for user in reaction.users(): + self.reactions[str(reaction)] += user.id + + def dict(self): + return self.__dict__ + + +class ChannelLog: + def __init__(self, channel: Union[discord.TextChannel, dict]): + if isinstance(channel, discord.TextChannel): + self.id = channel.id + self.name = channel.name + self.last_message_id = None + self.messages = [] + elif isinstance(channel, dict): + self.id = channel["id"] + self.name = channel["name"] + self.last_message_id = channel["last_message_id"] + self.messages = [MessageLog(message) for message in channel["messages"]] + + async def load(self, channel: discord.TextChannel): + self.name = channel.name + if self.last_message_id is not None: # append + while self.last_message_id != channel.last_message_id: + async for message in channel.history( + limit=CHUNK_SIZE, after=self.last_message_id, oldest_first=True + ): + self.last_message_id = message.id + m = MessageLog(message) + await m.load(message) + self.messages.insert(0, m) + yield len(self.messages), False + else: # first load + last_message_id = None + done = 0 + while done >= CHUNK_SIZE or last_message_id is None: + done = 0 + async for message in channel.history( + limit=CHUNK_SIZE, after=self.last_message_id, oldest_first=False + ): + done += 1 + last_message_id = message.id + m = MessageLog(message) + await m.load(message) + self.messages += [m] + yield len(self.messages), False + self.last_message_id == channel.last_message_id + yield len(self.messages), True + + def dict(self): + tmp = self.__dict__ + tmp["messages"] = [message.dict() for message in self.messages] + return tmp + + +class GuildLogs: + def __init__(self, guild: discord.Guild): + self.guild = guild + self.log_file = os.path.join(LOG_DIR, f"{guild}.logz") + self.channels = {} + + def dict(self): + return {id: self.channels[id].dict() for id in self.channels} + + async def load(self, target_channels: List[discord.TextChannel] = []): + # read logs + if os.path.exists(self.log_file): + channels = {} + with open(self.log_file, mode="r") as f: + channels = json.loads(f.readline().strip()) + self.channels = {id: ChannelLog(channels[id]) for id in channels} + # load channels + if len(target_channels) == 0: + target_channels = self.guild.text_channels + loading_new = False + total_msg = 0 + total_chan = 0 + for channel in target_channels: + if channel.id not in self.channels: + loading_new = True + self.channels[channel.id] = ChannelLog(channel) + async for count, done in self.channels[channel.id].load(channel): + yield ( + total_msg + count, + total_chan + (1 if done else 0), + loading_new, + False, + ) + total_msg += len(self.channels[channel.id].messages) + total_chan += 1 + yield total_msg, total_chan, loading_new, True + # write logs + with open(self.log_file, mode="w") as f: + f.write(json.dump(self.dict()))