From e02516677a5aa28526d3075ddf10032c34178e6a Mon Sep 17 00:00:00 2001 From: klemek Date: Wed, 13 Jan 2021 15:25:02 +0100 Subject: [PATCH] cancel command --- src/logs/__init__.py | 2 +- src/logs/guild_logs.py | 37 ++++++++++++++++++++++++++++++++++++- src/main.py | 7 +++++++ src/scanners/scanner.py | 11 ++++++++--- 4 files changed, 52 insertions(+), 5 deletions(-) diff --git a/src/logs/__init__.py b/src/logs/__init__.py index dbed38a..358e9af 100644 --- a/src/logs/__init__.py +++ b/src/logs/__init__.py @@ -1,3 +1,3 @@ from .message_log import MessageLog from .channel_logs import ChannelLogs -from .guild_logs import GuildLogs +from .guild_logs import GuildLogs, ALREADY_RUNNING, CANCELLED diff --git a/src/logs/guild_logs.py b/src/logs/guild_logs.py index 84472c8..229a682 100644 --- a/src/logs/guild_logs.py +++ b/src/logs/guild_logs.py @@ -16,6 +16,10 @@ LOG_DIR = "logs" current_analysis = [] +ALREADY_RUNNING = -100 +CANCELLED = -200 + + class GuildLogs: def __init__(self, guild: discord.Guild): self.guild = guild @@ -25,6 +29,9 @@ class GuildLogs: def dict(self) -> dict: return {id: self.channels[id].dict() for id in self.channels} + def check_cancelled(self) -> bool: + return self.log_file not in current_analysis + async def load( self, progress: discord.Message, @@ -34,7 +41,7 @@ class GuildLogs: ) -> Tuple[int, int]: global current_analysis if self.log_file in current_analysis: - return -1, -1 + return ALREADY_RUNNING, 0 current_analysis += [self.log_file] t00 = datetime.now() # read logs @@ -49,16 +56,22 @@ class GuildLogs: with open(self.log_file, mode="rb") as f: gziped_data = f.read() logging.info(f"log {self.guild.id} > read in {delta(t0):,}ms") + if self.check_cancelled(): + return CANCELLED, 0 await code_message(progress, "Reading saved history (2/4)...") t0 = datetime.now() json_data = gzip.decompress(gziped_data) logging.info( f"log {self.guild.id} > gzip decompress in {delta(t0):,}ms" ) + if self.check_cancelled(): + return CANCELLED, 0 await code_message(progress, "Reading saved history (3/4)...") t0 = datetime.now() channels = json.loads(json_data) logging.info(f"log {self.guild.id} > json parse in {delta(t0):,}ms") + if self.check_cancelled(): + return CANCELLED, 0 await code_message(progress, "Reading saved history (4/4)...") t0 = datetime.now() self.channels = {int(id): ChannelLogs(channels[id]) for id in channels} @@ -103,6 +116,8 @@ class GuildLogs: queried_msg = 0 total_chan = 0 max_chan = len(target_channels) + if self.check_cancelled(): + return CANCELLED, 0 await code_message( progress, f"Reading new history...\n0 messages in 0/{max_chan:,} channels\n(this might take a while)", @@ -123,6 +138,8 @@ class GuildLogs: warning_msg = ( "(some channels are new, this might take a long while)" ) + if self.check_cancelled(): + return CANCELLED, 0 await code_message( progress, f"Reading new history...\n{tmp_msg:,} messages in {total_chan + 1:,}/{max_chan:,} channels ({round(tmp_queried_msg/deltas(t0)):,}m/s)\n{warning_msg}", @@ -139,6 +156,8 @@ class GuildLogs: [len(channel.messages) for channel in self.channels.values()] ) real_total_chan = len(self.channels) + if self.check_cancelled(): + return CANCELLED, 0 await code_message( progress, f"Saving history (1/3)...\n{real_total_msg:,} messages in {real_total_chan:,} channels", @@ -148,6 +167,8 @@ class GuildLogs: logging.info( f"log {self.guild.id} > json dump in {delta(t0):,}ms -> {real_total_msg / deltas(t0):,.3f} m/s" ) + if self.check_cancelled(): + return CANCELLED, 0 await code_message( progress, f"Saving history (2/3)...\n{real_total_msg:,} messages in {real_total_chan:,} channels", @@ -157,6 +178,8 @@ class GuildLogs: logging.info( f"log {self.guild.id} > gzip in {delta(t0):,}ms -> {real_total_msg / deltas(t0):,.3f} m/s" ) + if self.check_cancelled(): + return CANCELLED, 0 await code_message( progress, f"Saving history (3/3)...\n{real_total_msg:,} messages in {real_total_chan:,} channels", @@ -167,6 +190,8 @@ class GuildLogs: logging.info( f"log {self.guild.id} > saved in {delta(t0):,}ms -> {real_total_msg / deltas(t0):,.3f} m/s" ) + if self.check_cancelled(): + return CANCELLED, 0 await code_message( progress, f"Analysing...\n{total_msg:,} messages in {total_chan:,} channels", @@ -174,3 +199,13 @@ class GuildLogs: logging.info(f"log {self.guild.id} > TOTAL TIME: {delta(t00):,}ms") current_analysis.remove(self.log_file) return total_msg, total_chan + + @staticmethod + async def cancel(client: discord.client, message: discord.Message, *args: str): + logs = GuildLogs(message.guild) + if logs.log_file in current_analysis: + current_analysis.remove(logs.log_file) + else: + await message.channel.send( + f"No analysis are currently running on this server", reference=message + ) diff --git a/src/main.py b/src/main.py index aff5ed6..743c3b7 100644 --- a/src/main.py +++ b/src/main.py @@ -9,6 +9,7 @@ from scanners import ( CompositionScanner, PresenceScanner, ) +from logs import GuildLogs logging.basicConfig( format="[%(asctime)s][%(levelname)s][%(module)s] %(message)s", level=logging.INFO @@ -24,6 +25,12 @@ bot = Bot( bot.log_calls = True +bot.register_command( + "(cancel|stop)", + GuildLogs.cancel, + "cancel: stop current analysis", + "", +) bot.register_command( "(emojis?|emotes?)", lambda *args: EmotesScanner().compute(*args), diff --git a/src/scanners/scanner.py b/src/scanners/scanner.py index 453fbd3..35c3a9f 100644 --- a/src/scanners/scanner.py +++ b/src/scanners/scanner.py @@ -6,7 +6,7 @@ import re import discord from utils import no_duplicate, get_intro, delta, deltas, mention, channel_mention -from logs import GuildLogs, ChannelLogs, MessageLog +from logs import GuildLogs, ChannelLogs, MessageLog, ALREADY_RUNNING, CANCELLED class Scanner(ABC): @@ -94,9 +94,14 @@ class Scanner(ABC): total_msg, total_chan = await logs.load( progress, self.channels, fast="fast" in args ) - if total_msg == -1: + if total_msg == CANCELLED: await message.channel.send( - f"An analysis is already running on this server, please be patient.", + "Operation cancelled by user", + reference=message, + ) + elif total_msg == ALREADY_RUNNING: + await message.channel.send( + "An analysis is already running on this server, please be patient.", reference=message, ) else: