From 84482969f132a7ea39e1633a4595261a85dd8dd5 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Thu, 26 Sep 2019 22:34:37 +0530 Subject: [PATCH] Introduce DownloadStatus class for storing stats of downloads Signed-off-by: lzzy12 --- bot/helper/bot_utils.py | 24 ++++++++++++++ bot/helper/download_status.py | 60 +++++++++++++++++++++++++++++++---- bot/helper/download_tools.py | 45 ++++++++++---------------- bot/helper/gdriveTools.py | 7 ++-- bot/helper/listeners.py | 8 ++--- bot/mirror.py | 40 +++++++++++++++-------- 6 files changed, 128 insertions(+), 56 deletions(-) create mode 100644 bot/helper/bot_utils.py diff --git a/bot/helper/bot_utils.py b/bot/helper/bot_utils.py new file mode 100644 index 0000000..aaefc25 --- /dev/null +++ b/bot/helper/bot_utils.py @@ -0,0 +1,24 @@ +from bot import download_list + + +def get_download(update_id): + return download_list[update_id].download() + + +def get_download_status_list(): + return list(download_list.values()) + + +def get_download_index(_list, gid): + index = 0 + for i in _list: + if i.download().gid == gid: + return index + index += 1 + + +def get_download_str_list(): + str_list = [] + for status in list(download_list.values()): + str_list.append(status.progress() + status.speed() + status.status()) + return str_list diff --git a/bot/helper/download_status.py b/bot/helper/download_status.py index 5427c2c..5278632 100644 --- a/bot/helper/download_status.py +++ b/bot/helper/download_status.py @@ -1,8 +1,54 @@ -class Message: - pass +from bot import aria2 -class Download: - def __init__(self, gid, download): - self.isDownloading = download.is_complete - self.gid = gid - self.download = download + +def get_download(gid): + return aria2.get_download(gid) + + +class DownloadStatus: + STATUS_UPLOADING = "Uploading" + STATUS_DOWNLOADING = "Downloading" + STATUS_WAITING = "Queued" + STATUS_FAILED = "Failed. Cleaning download" + + def __init__(self, gid): + self.__gid = gid + self.__download = get_download(gid) + + def __update(self): + self.__download = get_download(self.__gid) + + def progress(self): + self.__update() + return self.__download.progress_string() + + def speed(self): + self.__update() + return self.__download.download_speed_string() + + def name(self): + return self.__download.name + + def size(self): + return self.__download.total_length_string() + + def eta(self): + self.__update() + return self.__download.eta_string() + + def status(self): + self.__update() + status = None + if self.__download.is_waiting: + status = DownloadStatus.STATUS_WAITING + elif self.__download.is_complete: + # If download exists and is complete the it must be uploading + # otherwise the gid would have been removed from the download_list + status = DownloadStatus.STATUS_UPLOADING + elif self.__download.is_active: + status = DownloadStatus.STATUS_DOWNLOADING + return status + + def download(self): + self.__update() + return self.__download \ No newline at end of file diff --git a/bot/helper/download_tools.py b/bot/helper/download_tools.py index fd75f7b..7db2f72 100644 --- a/bot/helper/download_tools.py +++ b/bot/helper/download_tools.py @@ -1,7 +1,8 @@ from time import sleep from bot import LOGGER, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, download_list, aria2 from bot.helper.listeners import * -from urllib.parse import urlparse +from .download_status import DownloadStatus +from .bot_utils import * class DownloadHelper: @@ -27,34 +28,17 @@ class DownloadHelper: if self.is_url(link): download = aria2.add_uris([link], {'dir': DOWNLOAD_DIR + str(self.__listener.update.update_id)}) elif self.is_magnet(link): - download = aria2.add_magnet(link, {'dir': DOWNLOAD_DIR}) + download = aria2.add_magnet(link, {'dir': DOWNLOAD_DIR + str(self.__listener.update.update_id)}) self.__is_torrent = True else: self.__listener.onDownloadError("No download URL or URL malformed") return - download_list[self.__listener.update.update_id] = download.gid + download_list[self.__listener.update.update_id] = DownloadStatus(download.gid) self.__listener.onDownloadStarted(link) self.__update_download_status() - def get_downloads_status_str_list(self): - """ - Generates a human readable string of progress of all the downloads - :return: list of strings of progress of all the downloads - """ - str_list = [] - LOGGER.info(download_list) - for gid in list(download_list.values()): - download = aria2.get_download(gid) - str_list.append("" + download.name + ":- " - + download.progress_string() + " of " - + download.total_length_string() - + " at " + download.download_speed_string() - + " ,ETA: " + download.eta_string() - ) - return str_list - def __get_download(self): - return aria2.get_download(download_list[self.__listener.update.update_id]) + return get_download(self.__listener.update.update_id) def __get_followed_download_gid(self): download = self.__get_download() @@ -63,10 +47,6 @@ class DownloadHelper: return None def __update_download_status(self): - # TODO: Try to find a way to move this code to mirror.py and send only a - # list of Download objects to onDownloadProgress() - previous = None - LOGGER.info(self.get_downloads_status_str_list()) if self.__is_torrent: # Waiting for the actual gid new_gid = None @@ -74,16 +54,23 @@ class DownloadHelper: # Check every few seconds sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL) new_gid = self.__get_followed_download_gid() - download_list[self.__listener.update.update_id] = new_gid + self.__listener.onDownloadProgress(get_download_status_list()) + download_list[self.__listener.update.update_id] = DownloadStatus(new_gid) + # Start tracking the actual download + previous = None + status_list = get_download_status_list() while not self.__get_download().is_complete: if self.__get_download().has_failed: self.__listener.onDownloadError(self.__get_download().error_message) break - progress_str_list = self.get_downloads_status_str_list() + # TODO: Find a better way + progress_str_list = get_download_str_list() if progress_str_list != previous: - self.__listener.onDownloadProgress(progress_str_list) + self.__listener.onDownloadProgress(status_list, + get_download_index(status_list, self.__get_download().gid)) previous = progress_str_list + sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL) - self.__listener.onDownloadComplete(self.__get_download()) + self.__listener.onDownloadComplete(status_list, get_download_index(status_list, self.__get_download().gid)) diff --git a/bot/helper/gdriveTools.py b/bot/helper/gdriveTools.py index 1f6a0e6..94aad5a 100644 --- a/bot/helper/gdriveTools.py +++ b/bot/helper/gdriveTools.py @@ -8,6 +8,7 @@ import os from bot import LOGGER, parent_id, DOWNLOAD_DIR, download_list from .listeners import MirrorListeners from .fs_utils import clean_download +from .bot_utils import * class GoogleDriveHelper: @@ -59,7 +60,9 @@ class GoogleDriveHelper: return file_name, mime_type def upload(self, file_name: str): - self.__listener.onUploadStarted(file_name) + _list = get_download_status_list() + index = get_download_index(_list, get_download(self.__listener.update.update_id).gid) + self.__listener.onUploadStarted(_list, index) file_dir = "{}{}".format(DOWNLOAD_DIR, self.__listener.update.update_id) file_path = "{}/{}".format(file_dir, file_name) link = None @@ -87,7 +90,7 @@ class GoogleDriveHelper: raise Exception('Error: {}'.format(str(e))) del download_list[self.__listener.update.update_id] LOGGER.info(download_list) - self.__listener.onUploadComplete(link, file_name) + self.__listener.onUploadComplete(link, _list, index) LOGGER.info("Deleting downloaded file/folder..") clean_download(file_dir) return link diff --git a/bot/helper/listeners.py b/bot/helper/listeners.py index b6c983f..e6a7c39 100644 --- a/bot/helper/listeners.py +++ b/bot/helper/listeners.py @@ -7,19 +7,19 @@ class MirrorListeners: def onDownloadStarted(self, link: str): raise NotImplementedError - def onDownloadProgress(self, progress_str_list: list): + def onDownloadProgress(self, progress_status_list: list, index: int): raise NotImplementedError - def onDownloadComplete(self, download): + def onDownloadComplete(self, progress_status_list: list, index: int): raise NotImplementedError def onDownloadError(self, error: str): raise NotImplementedError - def onUploadStarted(self, filename: str): + def onUploadStarted(self, progress_status_list: list, index: int): raise NotImplementedError - def onUploadComplete(self, link: str, file_name: str): + def onUploadComplete(self, link: str, progress_status_list: list, index: int): raise NotImplementedError def onUploadError(self, error: str): diff --git a/bot/mirror.py b/bot/mirror.py index 06e7f88..c0b49cd 100644 --- a/bot/mirror.py +++ b/bot/mirror.py @@ -5,6 +5,19 @@ from bot import config, LOGGER, dispatcher LOGGER.info('mirror.py') +def get_readable_message(progress_list: list): + msg = "" + for status in progress_list: + msg += "Name: {}\n" \ + "status: {}\n" \ + "Downloaded: {} of {}\n" \ + "Speed: {}\n" \ + "ETA: {}\n\n".format(status.name(), status.status(), + status.progress(), status.size(), + status.speed(), status.eta()) + return msg + + class MirrorListener(listeners.MirrorListeners): def __init__(self, context, update, reply_message): super().__init__(context, update, reply_message) @@ -12,23 +25,21 @@ class MirrorListener(listeners.MirrorListeners): def onDownloadStarted(self, link): LOGGER.info("Adding link: " + link) - def onDownloadProgress(self, progress_str_list: list): + def onDownloadProgress(self, progress_status_list: list, index: int): LOGGER.info("Editing message") - msg = "Status:\n" - for progress_str in progress_str_list: - msg += progress_str + '\n\n' + msg = get_readable_message(progress_status_list) self.context.bot.edit_message_text(text=msg, message_id=self.reply_message.message_id, chat_id=self.reply_message.chat.id, parse_mode='HTMl') - def onDownloadComplete(self, download): - msg = "Downloaded" + def onDownloadComplete(self, progress_status_list, index: int): + msg = get_readable_message(progress_status_list) LOGGER.info("Download completed") self.context.bot.edit_message_text(text=msg, message_id=self.reply_message.message_id, chat_id=self.reply_message.chat.id, parse_mode='HTMl') gdrive = gdriveTools.GoogleDriveHelper(self) - gdrive.upload(download.name) + gdrive.upload(progress_status_list[index].name()) def onDownloadError(self, error): LOGGER.error(error) @@ -36,14 +47,15 @@ class MirrorListener(listeners.MirrorListeners): chat_id=self.reply_message.chat.id, parse_mode='HTMl') - def onUploadStarted(self, filename: str): - msg = "" + filename + ":- Uploading." - self.context.bot.edit_message_text(text=msg, message_id=self.reply_message.message_id, - chat_id=self.reply_message.chat.id, - parse_mode='HTMl') + def onUploadStarted(self, progress_status_list: list, index: int): + msg = get_readable_message(progress_status_list) + if not msg != self.update.message.text: + self.context.bot.edit_message_text(text=msg, message_id=self.reply_message.message_id, + chat_id=self.reply_message.chat.id, + parse_mode='HTMl') - def onUploadComplete(self, link: str, file_name: str): - msg = '{}'.format(link, file_name) + def onUploadComplete(self, link: str, progress_status_list: list, index: int): + msg = '{}'.format(link, progress_status_list[index].name()) self.context.bot.delete_message(chat_id=self.reply_message.chat.id, message_id=self.reply_message.message_id) self.context.bot.send_message(chat_id=self.update.message.chat_id, reply_to_message_id=self.update.message.message_id,