From 82ae49a33a861edbd386e2fee188850b9f721c5d Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Tue, 24 Dec 2019 17:24:35 +0530 Subject: [PATCH] Re implement upload progress after rebase Signed-off-by: lzzy12 --- bot/__init__.py | 7 +- bot/helper/ext_utils/bot_utils.py | 47 ++++++----- bot/helper/ext_utils/fs_utils.py | 2 +- .../mirror_utils/download_utils/__init__.py | 0 .../{ => download_utils}/aria2_download.py | 22 ++--- .../{ => download_utils}/download_helper.py | 0 .../mirror_utils/status_utils/__init__.py | 0 .../aria_download_status.py} | 34 ++------ .../status_utils/download_status.py | 36 ++++++++ .../{ => status_utils}/listeners.py | 10 +-- .../mirror_utils/status_utils/tar_status.py | 31 +++++++ .../status_utils/upload_status.py | 47 +++++++++++ .../mirror_utils/upload_utils/__init__.py | 0 .../{ => upload_utils}/gdriveTools.py | 51 ++++-------- bot/helper/telegram_helper/message_utils.py | 28 ++++--- bot/modules/list.py | 2 +- bot/modules/mirror.py | 83 ++++++++++--------- 17 files changed, 243 insertions(+), 157 deletions(-) create mode 100644 bot/helper/mirror_utils/download_utils/__init__.py rename bot/helper/mirror_utils/{ => download_utils}/aria2_download.py (75%) rename bot/helper/mirror_utils/{ => download_utils}/download_helper.py (100%) create mode 100644 bot/helper/mirror_utils/status_utils/__init__.py rename bot/helper/mirror_utils/{download_status.py => status_utils/aria_download_status.py} (62%) create mode 100644 bot/helper/mirror_utils/status_utils/download_status.py rename bot/helper/mirror_utils/{ => status_utils}/listeners.py (64%) create mode 100644 bot/helper/mirror_utils/status_utils/tar_status.py create mode 100644 bot/helper/mirror_utils/status_utils/upload_status.py create mode 100644 bot/helper/mirror_utils/upload_utils/__init__.py rename bot/helper/mirror_utils/{ => upload_utils}/gdriveTools.py (83%) diff --git a/bot/__init__.py b/bot/__init__.py index 19b45d4..0bd892a 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -3,7 +3,6 @@ import aria2p import threading import os from dotenv import load_dotenv -from telegram import Bot import telegram.ext as tg import time @@ -20,6 +19,7 @@ load_dotenv('config.env') Interval = [] + def getConfig(name: str): return os.environ[name] @@ -41,7 +41,6 @@ aria2 = aria2p.API( ) ) - DOWNLOAD_DIR = None BOT_TOKEN = None @@ -74,8 +73,6 @@ except KeyError as e: LOGGER.error("One or more env variables missing! Exiting now") exit(1) -bot = Bot(BOT_TOKEN) updater = tg.Updater(token=BOT_TOKEN) +bot = updater.bot dispatcher = updater.dispatcher - - diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index 713a6aa..c50651b 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/bot/helper/ext_utils/bot_utils.py @@ -3,7 +3,6 @@ import logging import re import threading import time -import math LOGGER = logging.getLogger(__name__) @@ -11,6 +10,7 @@ MAGNET_REGEX = r"magnet:\?xt=urn:btih:[a-zA-Z0-9]*" URL_REGEX = r"(?:(?:https?|ftp):\/\/)?[\w/\-?=%.]+\.[\w/\-?=%.]+" + class MirrorStatus: STATUS_UPLOADING = "Uploading" STATUS_DOWNLOADING = "Downloading" @@ -25,23 +25,25 @@ PROGRESS_INCOMPLETE = ['▏', '▎', '▍', '▌', '▋', '▊', '▉'] SIZE_UNITS = ['B', 'KB', 'MB', 'GB', 'TB', 'PB'] -class setInterval : - def __init__(self,interval,action) : - self.interval=interval - self.action=action - self.stopEvent=threading.Event() - thread=threading.Thread(target=self.__setInterval) + +class setInterval: + def __init__(self, interval, action): + self.interval = interval + self.action = action + self.stopEvent = threading.Event() + thread = threading.Thread(target=self.__setInterval) thread.start() - def __setInterval(self) : - nextTime=time.time()+self.interval - while not self.stopEvent.wait(nextTime-time.time()) : - nextTime+=self.interval + def __setInterval(self): + nextTime = time.time() + self.interval + while not self.stopEvent.wait(nextTime - time.time()): + nextTime += self.interval self.action() - def cancel(self) : + def cancel(self): self.stopEvent.set() + def get_readable_file_size(size_in_bytes) -> str: index = 0 while size_in_bytes >= 1024: @@ -65,10 +67,10 @@ def get_download_status_list(): def get_progress_bar_string(status): if status.status() == MirrorStatus.STATUS_UPLOADING: - completed = status.upload_helper.uploaded_bytes / 8 + completed = status.obj.uploaded_bytes / 8 else: completed = status.download().completed_length / 8 - total = status.download().total_length / 8 + total = status.size_raw() / 8 if total == 0: p = 0 else: @@ -99,6 +101,7 @@ def get_download_str(): result += (status.progress() + status.speed() + status.status()) return result + def get_readable_message(): with download_dict_lock: msg = "" @@ -110,11 +113,15 @@ def get_readable_message(): msg += "Archiving\n\n" elif download.status() == MirrorStatus.STATUS_WAITING: msg += "Queued\n\n" - else: - msg += f"{get_progress_bar_string(download)} {download.progress()} of {get_readable_file_size(download.download().total_length)}" \ - f" at {get_readable_file_size(download.download().download_speed)}ps, ETA: {download.eta()}\n\n" + elif download.status() == MirrorStatus.STATUS_DOWNLOADING: + msg += "Downloading" + if download.status() != MirrorStatus.STATUS_ARCHIVING: + msg += f"{get_progress_bar_string(download)} {download.progress()} of " \ + f"{download.size()}" \ + f" at {download.speed()}ps, ETA: {download.eta()}\n\n" return msg + def get_readable_time(seconds: int) -> str: result = '' (days, remainder) = divmod(seconds, 86400) @@ -135,14 +142,14 @@ def get_readable_time(seconds: int) -> str: def is_url(url: str): - url = re.findall(URL_REGEX,url) + url = re.findall(URL_REGEX, url) if url: return True - return False + return False def is_magnet(url: str): - magnet = re.findall(MAGNET_REGEX,url) + magnet = re.findall(MAGNET_REGEX, url) if magnet: return True return False diff --git a/bot/helper/ext_utils/fs_utils.py b/bot/helper/ext_utils/fs_utils.py index 54f8366..ad3b40f 100644 --- a/bot/helper/ext_utils/fs_utils.py +++ b/bot/helper/ext_utils/fs_utils.py @@ -5,6 +5,7 @@ import os import pathlib import mimetypes + def clean_download(path: str): if os.path.exists(path): LOGGER.info(f"Cleaning download: {path}") @@ -41,4 +42,3 @@ def get_mime_type(file_path): mime_type = mimetypes.guess_type(file_path)[0] mime_type = mime_type if mime_type else "text/plain" return mime_type - diff --git a/bot/helper/mirror_utils/download_utils/__init__.py b/bot/helper/mirror_utils/download_utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/bot/helper/mirror_utils/aria2_download.py b/bot/helper/mirror_utils/download_utils/aria2_download.py similarity index 75% rename from bot/helper/mirror_utils/aria2_download.py rename to bot/helper/mirror_utils/download_utils/aria2_download.py index d534a51..f04d012 100644 --- a/bot/helper/mirror_utils/aria2_download.py +++ b/bot/helper/mirror_utils/download_utils/aria2_download.py @@ -1,11 +1,12 @@ -from bot import DOWNLOAD_STATUS_UPDATE_INTERVAL, aria2, Interval +from bot import aria2 from bot.helper.ext_utils.bot_utils import * from .download_helper import DownloadHelper -from .download_status import DownloadStatus +from bot.helper.mirror_utils.status_utils.aria_download_status import AriaDownloadStatus from bot.helper.telegram_helper.message_utils import * import threading from aria2p import API + class AriaDownloadHelper(DownloadHelper): def __init__(self, listener): @@ -26,11 +27,11 @@ class AriaDownloadHelper(DownloadHelper): if self.gid == gid: if api.get_download(gid).followed_by_ids: self.gid = api.get_download(gid).followed_by_ids[0] - download_dict[self._listener.uid] = DownloadStatus(self.gid,self._listener) + download_dict[self._listener.uid] = AriaDownloadStatus(self.gid, self._listener) update_all_messages() LOGGER.info(f'Changed gid from {gid} to {self.gid}') else: - self._listener.onDownloadComplete() + self._listener.onDownloadComplete() def __onDownloadPause(self, api, gid): if self.gid == gid: @@ -56,14 +57,13 @@ class AriaDownloadHelper(DownloadHelper): download = aria2.add_uris([link], {'dir': path}) self.gid = download.gid with download_dict_lock: - download_dict[self._listener.uid] = DownloadStatus(self.gid,self._listener) + download_dict[self._listener.uid] = AriaDownloadStatus(self.gid, self._listener) if download.error_message: self._listener.onDownloadError(download.error_message) - return + return LOGGER.info(f"Started: {self.gid} DIR:{download.dir} ") aria2.listen_to_notifications(threaded=True, on_download_start=self.__onDownloadStarted, - on_download_error=self.__onDownloadError, - on_download_pause=self.__onDownloadPause, - on_download_stop=self.__onDownloadStopped, - on_download_complete=self.__onDownloadComplete) - + on_download_error=self.__onDownloadError, + on_download_pause=self.__onDownloadPause, + on_download_stop=self.__onDownloadStopped, + on_download_complete=self.__onDownloadComplete) diff --git a/bot/helper/mirror_utils/download_helper.py b/bot/helper/mirror_utils/download_utils/download_helper.py similarity index 100% rename from bot/helper/mirror_utils/download_helper.py rename to bot/helper/mirror_utils/download_utils/download_helper.py diff --git a/bot/helper/mirror_utils/status_utils/__init__.py b/bot/helper/mirror_utils/status_utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/bot/helper/mirror_utils/download_status.py b/bot/helper/mirror_utils/status_utils/aria_download_status.py similarity index 62% rename from bot/helper/mirror_utils/download_status.py rename to bot/helper/mirror_utils/status_utils/aria_download_status.py index f8e45ce..ba0948e 100644 --- a/bot/helper/mirror_utils/download_status.py +++ b/bot/helper/mirror_utils/status_utils/aria_download_status.py @@ -1,21 +1,22 @@ from bot import aria2, DOWNLOAD_DIR from bot.helper.ext_utils.bot_utils import get_readable_file_size, MirrorStatus, get_readable_time +from .download_status import DownloadStatus def get_download(gid): return aria2.get_download(gid) -class DownloadStatus: +class AriaDownloadStatus(DownloadStatus): def __init__(self, gid, listener): + super().__init__() self.upload_name = None self.is_archiving = False self.__gid = gid self.__download = get_download(gid) self.__uid = listener.uid self._listener = listener - self.upload_helper = None self.last = None self.is_waiting = False @@ -28,35 +29,20 @@ class DownloadStatus: :return: returns progress in percentage """ self.__update() - if self.upload_helper is not None: - return f'{round(self.upload_progress(), 2)}%' return self.__download.progress_string() - def upload_progress(self): - return self.upload_helper.uploaded_bytes / self.download().total_length * 100 - - def __size(self): + def size_raw(self): """ Gets total size of the mirror file/folder :return: total size of mirror """ return self.download().total_length - def __upload_speed(self): - """ - :return: Upload speed in Bytes/Seconds - """ - return self.upload_helper.speed() - def speed(self): self.__update() - if self.upload_helper is not None: - return f'{get_readable_file_size(self.__upload_speed())}/s' return self.__download.download_speed_string() def name(self): - if self.upload_name is not None: - return self.upload_name self.__update() return self.__download.name @@ -68,24 +54,14 @@ class DownloadStatus: def eta(self): self.__update() - if self.upload_helper is not None: - try: - seconds = (self.__size() - self.upload_helper.uploaded_bytes) / self.__upload_speed() - return f'{get_readable_time(seconds)}' - except ZeroDivisionError: - return '-' return self.__download.eta_string() def status(self): self.__update() - if self.is_archiving: - status = MirrorStatus.STATUS_ARCHIVING - elif self.download().is_waiting: + if self.download().is_waiting: status = MirrorStatus.STATUS_WAITING elif self.download().is_paused: status = MirrorStatus.STATUS_CANCELLED - elif self.upload_helper is not None: - status = MirrorStatus.STATUS_UPLOADING elif self.__download.has_failed: status = MirrorStatus.STATUS_FAILED else: diff --git a/bot/helper/mirror_utils/status_utils/download_status.py b/bot/helper/mirror_utils/status_utils/download_status.py new file mode 100644 index 0000000..4ebc214 --- /dev/null +++ b/bot/helper/mirror_utils/status_utils/download_status.py @@ -0,0 +1,36 @@ +# Generic status class. All other status classes must inherit this class + + +class DownloadStatus: + + def progress(self): + """ + Calculates the progress of the mirror (upload or download) + :return: progress in percentage + """ + raise NotImplementedError + + def speed(self): + """:return: speed in bytes per second""" + raise NotImplementedError + + def name(self): + """:return name of file/directory being processed""" + raise NotImplementedError + + def path(self): + """:return path of the file/directory""" + raise NotImplementedError + + def size(self): + """:return Size of file folder""" + raise NotImplementedError + + def eta(self): + """:return ETA of the process to complete""" + raise NotImplementedError + + def status(self): + """:return String describing what is the object of this class will be tracking (upload/download/something + else) """ + raise NotImplementedError diff --git a/bot/helper/mirror_utils/listeners.py b/bot/helper/mirror_utils/status_utils/listeners.py similarity index 64% rename from bot/helper/mirror_utils/listeners.py rename to bot/helper/mirror_utils/status_utils/listeners.py index 134cbe2..9c824b7 100644 --- a/bot/helper/mirror_utils/listeners.py +++ b/bot/helper/mirror_utils/status_utils/listeners.py @@ -17,14 +17,14 @@ class MirrorListeners: def onDownloadError(self, error: str): raise NotImplementedError - def onUploadStarted(self, progress_status_list: list, index: int): + def onUploadStarted(self): raise NotImplementedError - def onUploadProgress(self, progress: list, index: int): + def onUploadProgress(self): raise NotImplementedError - def onUploadComplete(self, link: str, progress_status_list: list, index: int): + def onUploadComplete(self, link: str): raise NotImplementedError - def onUploadError(self, error: str, progress_status: list, index: int): - raise NotImplementedError \ No newline at end of file + def onUploadError(self, error: str): + raise NotImplementedError diff --git a/bot/helper/mirror_utils/status_utils/tar_status.py b/bot/helper/mirror_utils/status_utils/tar_status.py new file mode 100644 index 0000000..5a21c3a --- /dev/null +++ b/bot/helper/mirror_utils/status_utils/tar_status.py @@ -0,0 +1,31 @@ +from .download_status import DownloadStatus +from bot.helper.ext_utils.bot_utils import get_readable_file_size, MirrorStatus + + +class TarStatus(DownloadStatus): + def __init__(self, name, path, size): + self.__name = name + self.__path = path + self.__size = size + + def progress(self): + return '0' + + def speed(self): + return '0' + + def name(self): + return self.__name + + def path(self): + return self.__path + + def size(self): + return get_readable_file_size(self.__size) + + def eta(self): + return '0s' + + def status(self): + return MirrorStatus.STATUS_ARCHIVING + diff --git a/bot/helper/mirror_utils/status_utils/upload_status.py b/bot/helper/mirror_utils/status_utils/upload_status.py new file mode 100644 index 0000000..df96a80 --- /dev/null +++ b/bot/helper/mirror_utils/status_utils/upload_status.py @@ -0,0 +1,47 @@ +from .download_status import DownloadStatus +from bot.helper.ext_utils.bot_utils import MirrorStatus, get_readable_file_size, get_readable_time +from bot import DOWNLOAD_DIR + + +class UploadStatus(DownloadStatus): + def __init__(self, obj, size, uid): + self.obj = obj + self.__size = size + self.uid = uid + + def path(self): + return f"{DOWNLOAD_DIR}{self.uid}" + + def size_raw(self): + return self.__size + + def size(self): + return get_readable_file_size(self.__size) + + def status(self): + return MirrorStatus.STATUS_UPLOADING + + def name(self): + return self.obj.name + + def progress_raw(self): + return self.obj.uploaded_bytes / self.__size * 100 + + def progress(self): + return f'{round(self.progress_raw(), 2)}%' + + def speed_raw(self): + """ + :return: Upload speed in Bytes/Seconds + """ + return self.obj.speed() + + def speed(self): + return f'{get_readable_file_size(self.speed_raw())}/s' + + def eta(self): + try: + seconds = (self.__size - self.obj.uploaded_bytes) / self.speed_raw() + return f'{get_readable_time(seconds)}' + except ZeroDivisionError: + return '-' diff --git a/bot/helper/mirror_utils/upload_utils/__init__.py b/bot/helper/mirror_utils/upload_utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/bot/helper/mirror_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py similarity index 83% rename from bot/helper/mirror_utils/gdriveTools.py rename to bot/helper/mirror_utils/upload_utils/gdriveTools.py index d135f56..bc81b3a 100644 --- a/bot/helper/mirror_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -4,19 +4,16 @@ from google.auth.transport.requests import Request from googleapiclient.http import MediaFileUpload import pickle import os -import time from bot import LOGGER, parent_id, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL from bot.helper.ext_utils.fs_utils import get_mime_type from bot.helper.ext_utils.bot_utils import * -from bot.helper.ext_utils.exceptions import MessageDeletedError -import threading logging.getLogger('googleapiclient.discovery').setLevel(logging.ERROR) class GoogleDriveHelper: - def __init__(self, listener=None): + def __init__(self, name=None, listener=None): self.__G_DRIVE_TOKEN_FILE = "token.pickle" # Check https://developers.google.com/drive/scopes for all available scopes self.__OAUTH_SCOPE = "https://www.googleapis.com/auth/drive.file" @@ -34,6 +31,8 @@ class GoogleDriveHelper: self.is_uploading = True self.is_cancelled = False self.status = None + self.updater = None + self.name = name def cancel(self): self.is_cancelled = True @@ -50,28 +49,12 @@ class GoogleDriveHelper: return 0 def _on_upload_progress(self): - while self.is_uploading: - if self.status is not None: - chunk_size = self.status.total_size * self.status.progress() - self._file_uploaded_bytes - self._file_uploaded_bytes = self.status.total_size * self.status.progress() - LOGGER.info(f'Chunk size: {get_readable_file_size(chunk_size)}') - self.uploaded_bytes += chunk_size - self.total_time += DOWNLOAD_STATUS_UPDATE_INTERVAL - - if self._should_update: - try: - LOGGER.info('Updating messages') - _list = get_download_status_list() - index = get_download_index(_list, get_download(self.__listener.message.message_id).gid) - self.__listener.onUploadProgress(_list, index) - except MessageDeletedError as e: - LOGGER.info(f'Stopped calling onDownloadProgress(): {str(e)}') - # TODO: Find a way to know if the Error is actually about message not found and not found - # self._should_update = False - pass - else: - LOGGER.info('status: None') - time.sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL) + if self.status is not None: + chunk_size = self.status.total_size * self.status.progress() - self._file_uploaded_bytes + self._file_uploaded_bytes = self.status.total_size * self.status.progress() + LOGGER.info(f'Chunk size: {get_readable_file_size(chunk_size)}') + self.uploaded_bytes += chunk_size + self.total_time += DOWNLOAD_STATUS_UPDATE_INTERVAL def __upload_empty_file(self, path, file_name, mime_type, parent_id=None): media_body = MediaFileUpload(path, @@ -135,14 +118,12 @@ class GoogleDriveHelper: return download_url def upload(self, file_name: str): - _list = get_download_status_list() - index = get_download_index(_list, get_download(self.__listener.message.message_id).gid) - self.__listener.onUploadStarted(_list, index) + self.__listener.onUploadStarted() file_dir = f"{DOWNLOAD_DIR}{self.__listener.message.message_id}" file_path = f"{file_dir}/{file_name}" LOGGER.info("Uploading File: " + file_path) self.start_time = time.time() - #threading.Thread(target=self._on_upload_progress).start() + self.updater = setInterval(1, self._on_upload_progress) if os.path.isfile(file_path): try: mime_type = get_mime_type(file_path) @@ -154,10 +135,10 @@ class GoogleDriveHelper: LOGGER.error(str(e)) e_str = str(e).replace('<', '') e_str = e_str.replace('>', '') - self.__listener.onUploadError(e_str, _list, index) + self.__listener.onUploadError(e_str) return finally: - self.is_uploading = False + self.updater.cancel() else: try: dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_id) @@ -170,12 +151,12 @@ class GoogleDriveHelper: LOGGER.error(str(e)) e_str = str(e).replace('<', '') e_str = e_str.replace('>', '') - self.__listener.onUploadError(e_str, _list, index) + self.__listener.onUploadError(e_str) return finally: - self.is_uploading = False + self.updater.cancel() LOGGER.info(download_dict) - self.__listener.onUploadComplete(link, _list, index) + self.__listener.onUploadComplete(link) LOGGER.info("Deleting downloaded file/folder..") return link diff --git a/bot/helper/telegram_helper/message_utils.py b/bot/helper/telegram_helper/message_utils.py index de758f0..ad4e6fe 100644 --- a/bot/helper/telegram_helper/message_utils.py +++ b/bot/helper/telegram_helper/message_utils.py @@ -7,10 +7,11 @@ from bot.helper.ext_utils.bot_utils import get_readable_message from telegram.error import TimedOut, BadRequest from bot import bot + def sendMessage(text: str, bot, update: Update): return bot.send_message(update.message.chat_id, - reply_to_message_id=update.message.message_id, - text=text, parse_mode='HTMl') + reply_to_message_id=update.message.message_id, + text=text, parse_mode='HTMl') def editMessage(text: str, message: Message): @@ -26,7 +27,7 @@ def editMessage(text: str, message: Message): def deleteMessage(bot, message: Message): try: bot.delete_message(chat_id=message.chat.id, - message_id=message.message_id) + message_id=message.message_id) except Exception as e: LOGGER.error(str(e)) @@ -34,8 +35,8 @@ def deleteMessage(bot, message: Message): def sendLogFile(bot, update: Update): with open('log.txt', 'rb') as f: bot.send_document(document=f, filename=f.name, - reply_to_message_id=update.message.message_id, - chat_id=update.message.chat_id) + reply_to_message_id=update.message.message_id, + chat_id=update.message.chat_id) def auto_delete_message(bot, cmd_message: Message, bot_message: Message): @@ -48,6 +49,7 @@ def auto_delete_message(bot, cmd_message: Message, bot_message: Message): except AttributeError: pass + def delete_all_messages(): with status_reply_dict_lock: for message in list(status_reply_dict.values()): @@ -59,28 +61,30 @@ def delete_all_messages(): del status_reply_dict[message.chat.id] pass + def update_all_messages(): msg = get_readable_message() with status_reply_dict_lock: for chat_id in list(status_reply_dict.keys()): if msg != status_reply_dict[chat_id].text: try: - editMessage(msg,status_reply_dict[chat_id]) - except: - pass + editMessage(msg, status_reply_dict[chat_id]) + except BadRequest as e: + LOGGER.error(str(e)) status_reply_dict[chat_id].text = msg -def sendStatusMessage(msg,bot): + +def sendStatusMessage(msg, bot): progress = get_readable_message() with status_reply_dict_lock: if msg.message.chat.id in list(status_reply_dict.keys()): try: message = status_reply_dict[msg.message.chat.id] - deleteMessage(bot,message) + deleteMessage(bot, message) del status_reply_dict[msg.message.chat.id] except Exception as e: LOGGER.error(str(e)) del status_reply_dict[msg.message.chat.id] pass - message = sendMessage(progress,bot,msg) - status_reply_dict[msg.message.chat.id] = message \ No newline at end of file + message = sendMessage(progress, bot, msg) + status_reply_dict[msg.message.chat.id] = message diff --git a/bot/modules/list.py b/bot/modules/list.py index 5f177ab..4969291 100644 --- a/bot/modules/list.py +++ b/bot/modules/list.py @@ -1,5 +1,5 @@ from telegram.ext import CommandHandler, run_async -from bot.helper.mirror_utils.gdriveTools import GoogleDriveHelper +from bot.helper.mirror_utils.upload_utils.gdriveTools import GoogleDriveHelper from bot import LOGGER, dispatcher from bot.helper.telegram_helper.message_utils import auto_delete_message, sendMessage import threading diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 74cce36..4f6193a 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -1,16 +1,17 @@ from telegram.ext import CommandHandler, run_async -from telegram.error import BadRequest, TimedOut -from bot.helper.mirror_utils import aria2_download, gdriveTools, listeners -from bot import LOGGER, dispatcher, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL +from bot.helper.mirror_utils.status_utils import listeners +from bot.helper.mirror_utils.upload_utils import gdriveTools +from bot.helper.mirror_utils.download_utils import aria2_download +from bot.helper.mirror_utils.status_utils.upload_status import UploadStatus +from bot.helper.mirror_utils.status_utils.tar_status import TarStatus +from bot import dispatcher, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL from bot.helper.ext_utils import fs_utils, bot_utils -from bot import download_dict, status_reply_dict, status_reply_dict_lock, download_dict_lock, Interval +from bot import Interval from bot.helper.telegram_helper.message_utils import * -from bot.helper.ext_utils.bot_utils import get_readable_message, MirrorStatus,setInterval -from bot.helper.ext_utils.exceptions import MessageDeletedError +from bot.helper.ext_utils.bot_utils import setInterval from bot.helper.telegram_helper.filters import CustomFilters from bot.helper.telegram_helper.bot_commands import BotCommands import pathlib -from time import sleep class MirrorListener(listeners.MirrorListeners): @@ -18,12 +19,13 @@ class MirrorListener(listeners.MirrorListeners): super().__init__(bot, update) self.isTar = isTar - def onDownloadStarted(self,reply_message): + def onDownloadStarted(self): pass def onDownloadProgress(self): # We are handling this on our own! pass + def clean(self): Interval[0].cancel() del Interval[0] @@ -33,12 +35,18 @@ class MirrorListener(listeners.MirrorListeners): with download_dict_lock: LOGGER.info(f"Download completed: {download_dict[self.uid].name()}") download = download_dict[self.uid] + name = download.name() + size = download.size_raw() + m_path = f'{DOWNLOAD_DIR}{self.uid}/{download.name()}' if self.isTar: download.is_archiving = True try: - path = fs_utils.tar(f'{DOWNLOAD_DIR}{self.uid}/{download.name()}') + with download_dict_lock: + download_dict[self.uid] = TarStatus(name, m_path, size) + path = fs_utils.tar(m_path) except FileNotFoundError: - self.onUploadError('Download cancelled!', download_dict, self.uid) + LOGGER.info('File to archive not found!') + self.onUploadError('Internal error occurred!!') return else: path = f'{DOWNLOAD_DIR}{self.uid}/{download_dict[self.uid].name()}' @@ -46,17 +54,17 @@ class MirrorListener(listeners.MirrorListeners): with download_dict_lock: download_dict[self.uid].is_archiving = False LOGGER.info(f"Upload Name : {name}") - download_dict[self.uid].upload_name = name - gdrive = gdriveTools.GoogleDriveHelper(self) - download_dict[self.uid].upload_helper = gdrive + drive = gdriveTools.GoogleDriveHelper(name, self) + upload_status = UploadStatus(drive, size, self.uid) + download_dict[self.uid] = upload_status update_all_messages() - gdrive.upload(name) + drive.upload(name) def onDownloadError(self, error): LOGGER.info(self.update.effective_chat.id) with download_dict_lock: try: - download = download_dict[self.uid] + download = download_dict[self.uid] del download_dict[self.uid] LOGGER.info(f"Deleting folder: {download.path()}") fs_utils.clean_download(download.path()) @@ -76,25 +84,32 @@ class MirrorListener(listeners.MirrorListeners): update_all_messages() sendMessage(msg, self.bot, self.update) - def onUploadStarted(self, progress_status_list: list, index: int): + def onUploadStarted(self): pass - def onUploadComplete(self, link: str, progress_status_list: list, index: int): - msg = f'{progress_status_list[index].name()} ({progress_status_list[index].size()})' + def onUploadComplete(self, link: str): with download_dict_lock: - del download_dict[self.uid] + msg = f'{download_dict[self.uid].name()} ({download_dict[self.uid].size()})' + LOGGER.info(f'Done Uploading {download_dict[self.uid].name()}') + if len(download_dict) == 0: self.clean() else: update_all_messages() sendMessage(msg, self.bot, self.update) - print("Downloads: "+str(download_dict)) try: - fs_utils.clean_download(progress_status_list[index].path()) + with download_dict_lock: + fs_utils.clean_download(download_dict[self.uid].path()) + del download_dict[self.uid] + count = len(download_dict) + if count == 0: + self.clean() + else: + update_all_messages() except FileNotFoundError: pass - def onUploadError(self, error: str, progress_status: list, index: int): + def onUploadError(self, error: str): LOGGER.error(error) sendMessage(error, self.bot, self.update) with download_dict_lock: @@ -104,22 +119,13 @@ class MirrorListener(listeners.MirrorListeners): else: update_all_messages() try: - fs_utils.clean_download(progress_status[index].path()) + with download_dict_lock: + fs_utils.clean_download(download_dict[self.uid].path()) except FileNotFoundError: pass - def onUploadProgress(self, progress: list, index: int): - msg = get_readable_message(progress) - try: - print("lul") - #editMessage(msg, self.bot, self.reply_message) - except BadRequest as e: - raise MessageDeletedError(str(e)) - except TimedOut: - pass - -def _mirror(bot,update ,isTar=False): +def _mirror(bot, update, isTar=False): message_args = update.message.text.split(' ') try: link = message_args[1] @@ -142,13 +148,14 @@ def _mirror(bot,update ,isTar=False): listener = MirrorListener(bot, update, isTar) aria = aria2_download.AriaDownloadHelper(listener) aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/') - sendStatusMessage(update,bot) + sendStatusMessage(update, bot) if len(Interval) == 0: - Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL,update_all_messages)) + Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages)) + @run_async -def mirror(bot,update): - _mirror(bot,update) +def mirror(bot, update): + _mirror(bot, update) @run_async