diff --git a/bot/__main__.py b/bot/__main__.py index 4478eec..44d0444 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -1,6 +1,6 @@ from telegram.ext import CommandHandler, run_async from bot import dispatcher, LOGGER, updater, aria2 -import bot.mirror,bot.list +import bot.mirror, bot.list, bot.mirror_status from bot.helper import fs_utils import signal diff --git a/bot/helper/bot_utils.py b/bot/helper/bot_utils.py index 2573057..c313afe 100644 --- a/bot/helper/bot_utils.py +++ b/bot/helper/bot_utils.py @@ -1,3 +1,4 @@ + from bot import download_dict @@ -35,3 +36,10 @@ def get_readable_message(progress_list: list = download_dict.values()): status.progress(), status.size(), status.speed(), status.eta()) return msg + + +# Custom Exception class for killing thread as soon as they aren't needed +class KillThreadException(Exception): + def __init__(self, message, error=None): + super().__init__(message) + self.error = error diff --git a/bot/helper/download_tools.py b/bot/helper/download_tools.py index ac74c45..7214e89 100644 --- a/bot/helper/download_tools.py +++ b/bot/helper/download_tools.py @@ -55,7 +55,10 @@ class DownloadHelper: # Check every few seconds sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL) new_gid = self.__get_followed_download_gid() - self.__listener.onDownloadProgress(get_download_status_list(), index) + try: + self.__listener.onDownloadProgress(get_download_status_list(), index) + except KillThreadException: + return download_dict[self.__listener.update.update_id] = DownloadStatus(new_gid, self.__listener.update.update_id) # Start tracking the actual download @@ -70,7 +73,10 @@ class DownloadHelper: # TODO: Find a better way to differentiate between 2 list of objects progress_str_list = get_download_str() if progress_str_list != previous: - self.__listener.onDownloadProgress(status_list, index) + try: + self.__listener.onDownloadProgress(status_list, index) + except KillThreadException: + return previous = progress_str_list download = self.__get_download() sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL) diff --git a/bot/helper/message_utils.py b/bot/helper/message_utils.py index 65ba56f..cc2b18c 100644 --- a/bot/helper/message_utils.py +++ b/bot/helper/message_utils.py @@ -1,7 +1,7 @@ from telegram.error import BadRequest from telegram.message import Message from telegram.update import Update - +from bot import LOGGER def sendMessage(text: str, context, update: Update): return context.bot.send_message(update.message.chat_id, @@ -10,12 +10,9 @@ def sendMessage(text: str, context, update: Update): def editMessage(text: str, context, message: Message): - try: context.bot.edit_message_text(text=text, message_id=message.message_id, chat_id=message.chat.id, parse_mode='HTMl') - except BadRequest: - pass def deleteMessage(context, message: Message): diff --git a/bot/mirror.py b/bot/mirror.py index 6a51045..25fdea1 100644 --- a/bot/mirror.py +++ b/bot/mirror.py @@ -5,7 +5,7 @@ from bot.helper import fs_utils from bot import download_dict, status_reply_dict, DOWNLOAD_STATUS_UPDATE_INTERVAL from bot.helper.message_utils import * from time import sleep -from bot.helper.bot_utils import get_readable_message +from bot.helper.bot_utils import get_readable_message, KillThreadException LOGGER.info('mirror.py') @@ -19,7 +19,10 @@ class MirrorListener(listeners.MirrorListeners): def onDownloadProgress(self, progress_status_list: list, index: int): msg = get_readable_message(progress_status_list) LOGGER.info("Editing message") - editMessage(msg, self.context, self.reply_message) + try: + editMessage(msg, self.context, self.reply_message) + except BadRequest: + raise KillThreadException('Message deleted. Terminate thread') def onDownloadComplete(self, progress_status_list, index: int): msg = get_readable_message(progress_status_list) @@ -70,28 +73,6 @@ def mirror(update, context): aria.add_download(link) -@run_async -def mirror_status(update: Update, context): - message = get_readable_message() - if len(message) == 0: - message = "No active downloads" - sendMessage(message, context, update) - return - index = update.effective_chat.id - if index in status_reply_dict.keys(): - deleteMessage(context, status_reply_dict[index]) - del status_reply_dict[index] - while True: - message = get_readable_message() - if index in status_reply_dict.keys(): - if len(message) == 0: - message = "No active downloads" - editMessage(message, context, status_reply_dict[index]) - break - editMessage(message, context, status_reply_dict[index]) - else: - status_reply_dict[index] = sendMessage(message, context, update) - sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL) @run_async @@ -100,6 +81,4 @@ def cancel_mirror(update, context): mirror_handler = CommandHandler('mirror', mirror) -mirror_status_handler = CommandHandler('status', mirror_status) dispatcher.add_handler(mirror_handler) -dispatcher.add_handler(mirror_status_handler) diff --git a/bot/mirror_status.py b/bot/mirror_status.py new file mode 100644 index 0000000..41729f8 --- /dev/null +++ b/bot/mirror_status.py @@ -0,0 +1,43 @@ +from telegram.ext import CommandHandler, run_async +from bot import dispatcher, status_reply_dict, DOWNLOAD_STATUS_UPDATE_INTERVAL +from bot.helper.message_utils import * +from time import sleep +from bot.helper.bot_utils import get_readable_message + + +@run_async +def mirror_status(update: Update, context): + message = get_readable_message() + if len(message) == 0: + message = "No active downloads" + sendMessage(message, context, update) + return + index = update.effective_chat.id + if index in status_reply_dict.keys(): + deleteMessage(context, status_reply_dict[index]) + del status_reply_dict[index] + kill_thread = False + while len(message) != 0: + message = get_readable_message() + if index in status_reply_dict.keys(): + if len(message) == 0: + message = "No active downloads" + editMessage(message, context, status_reply_dict[index]) + break + try: + editMessage(message, context, status_reply_dict[index]) + except BadRequest: + break + else: + # If the loop returns here 2nd time, it means the message + # has been replaced by a new message due to a second /status command in the chat. + # So we kill the thread by simply breaking the loop + if kill_thread: + break + status_reply_dict[index] = sendMessage(message, context, update) + kill_thread = True + sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL) + + +mirror_status_handler = CommandHandler('status', mirror_status) +dispatcher.add_handler(mirror_status_handler)