diff --git a/bot/helper/bot_utils.py b/bot/helper/bot_utils.py index 7e6f3fd..cbd4d4f 100644 --- a/bot/helper/bot_utils.py +++ b/bot/helper/bot_utils.py @@ -1,9 +1,31 @@ from bot import download_dict -from bot.helper.download_status import DownloadStatus + + +class MirrorStatus: + + STATUS_UPLOADING = "Uploading" + STATUS_DOWNLOADING = "Downloading" + STATUS_WAITING = "Queued" + STATUS_FAILED = "Failed. Cleaning download" + STATUS_CANCELLED = "Cancelled" + PROGRESS_MAX_SIZE = 100 // 8 PROGRESS_INCOMPLETE = ['▏', '▎', '▍', '▌', '▋', '▊', '▉'] +SIZE_UNITS = ['B', 'KB', 'MB', 'GB', 'TB', 'PB'] + + +def get_readable_file_size(size_in_bytes) -> str: + index = 0 + while size_in_bytes >= 1024: + size_in_bytes /= 1024 + index += 1 + try: + return f'{round(size_in_bytes, 2)} {SIZE_UNITS[index]}' + except IndexError: + return 'File too large' + def get_download(message_id): return download_dict[message_id].download() @@ -13,8 +35,11 @@ def get_download_status_list(): return list(download_dict.values()) -def get_progress_bar_string(status: DownloadStatus): - completed = status.download().completed_length/8 +def get_progress_bar_string(status): + if status.status() == MirrorStatus.STATUS_UPLOADING: + completed = status.uploaded_bytes/8 + else: + completed = status.download().completed_length/8 total = status.download().total_length/8 if total == 0: p = 0 @@ -50,18 +75,8 @@ def get_readable_message(progress_list: list = download_dict.values()): msg = '' for status in progress_list: msg += f'Name: {status.name()}\n' \ - f'status: {status.status()}\n' - - if status.status() == DownloadStatus.STATUS_DOWNLOADING: - msg += f'{get_progress_bar_string(status)} {status.progress()} of {status.size()}\n' \ + f'status: {status.status()}\n' \ + f'{get_progress_bar_string(status)} {status.progress()} of {status.size()}\n' \ f'Speed: {status.speed()}\n' \ f'ETA: {status.eta()}\n' - msg += '\n' - 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 + return msg \ No newline at end of file diff --git a/bot/helper/download_status.py b/bot/helper/download_status.py index 0487f72..069ca20 100644 --- a/bot/helper/download_status.py +++ b/bot/helper/download_status.py @@ -1,4 +1,5 @@ -from bot import aria2, download_dict, DOWNLOAD_DIR +from bot import aria2, DOWNLOAD_DIR +from .bot_utils import get_readable_file_size, MirrorStatus def get_download(gid): @@ -6,26 +7,51 @@ def get_download(gid): class DownloadStatus: - STATUS_UPLOADING = "Uploading" - STATUS_DOWNLOADING = "Downloading" - STATUS_WAITING = "Queued" - STATUS_FAILED = "Failed. Cleaning download" - STATUS_CANCELLED = "Cancelled" def __init__(self, gid, message_id): self.__gid = gid self.__download = get_download(gid) self.__uid = message_id + self.uploaded_bytes = 0 + self.upload_time = 0 def __update(self): self.__download = get_download(self.__gid) def progress(self): + """ + Calculates the progress of the mirror (upload or download) + :return: returns progress in percentage + """ self.__update() + if self.status() == MirrorStatus.STATUS_UPLOADING: + return f'{round(self.upload_progress(), 2)}%' return self.__download.progress_string() + def upload_progress(self): + return self.uploaded_bytes / self.download().total_length * 100 + + def __size(self): + """ + Gets total size of the mirror file/folder + :return: total size of mirror + """ + return self.download().total_length + + def __upload_speed(self): + """ + Calculates upload speed in bytes/second + :return: Upload speed in Bytes/Seconds + """ + try: + return self.uploaded_bytes / self.upload_time + except ZeroDivisionError: + return 0 + def speed(self): self.__update() + if self.status() == MirrorStatus.STATUS_UPLOADING: + return f'{get_readable_file_size(self.__upload_speed())}/s' return self.__download.download_speed_string() def name(self): @@ -39,23 +65,28 @@ class DownloadStatus: def eta(self): self.__update() + if self.status() == MirrorStatus.STATUS_UPLOADING: + try: + return f'{round(self.__size() / self.__upload_speed(), 2)} seconds' + except ZeroDivisionError: + return '-' return self.__download.eta_string() def status(self): self.__update() status = None if self.__download.is_waiting: - status = DownloadStatus.STATUS_WAITING + status = MirrorStatus.STATUS_WAITING elif self.download().is_paused: - status = DownloadStatus.STATUS_CANCELLED + status = MirrorStatus.STATUS_CANCELLED 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 + status = MirrorStatus.STATUS_UPLOADING elif self.__download.has_failed: - status = DownloadStatus.STATUS_FAILED + status = MirrorStatus.STATUS_FAILED elif self.__download.is_active: - status = DownloadStatus.STATUS_DOWNLOADING + status = MirrorStatus.STATUS_DOWNLOADING return status def download(self): diff --git a/bot/helper/exceptions.py b/bot/helper/exceptions.py index 3affc67..0710922 100644 --- a/bot/helper/exceptions.py +++ b/bot/helper/exceptions.py @@ -1,2 +1,8 @@ class DriveAuthError(Exception): - pass \ No newline at end of file + pass + +# 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/gdriveTools.py b/bot/helper/gdriveTools.py index 163ef93..f741994 100644 --- a/bot/helper/gdriveTools.py +++ b/bot/helper/gdriveTools.py @@ -7,7 +7,7 @@ import os from bot import LOGGER, parent_id, DOWNLOAD_DIR from .fs_utils import get_mime_type from .bot_utils import * - +import time import logging logging.getLogger('googleapiclient.discovery').setLevel(logging.ERROR) @@ -25,6 +25,8 @@ class GoogleDriveHelper: self.__G_DRIVE_BASE_DOWNLOAD_URL = "https://drive.google.com/uc?id={}&export=download" self.__listener = listener self.__service = self.authorize() + self.uploadedBytes = 0 + self.start_time = 0 def upload_file(self, file_path, file_name, mime_type, parent_id): # File body description @@ -47,11 +49,27 @@ class GoogleDriveHelper: 'withLink': True } # Insert a file - drive_file = self.__service.files().create(body=file_metadata, media_body=media_body).execute() + drive_file = self.__service.files().create(body=file_metadata, media_body=media_body) + response = None + _list = get_download_status_list() + index = get_download_index(_list, get_download(self.__listener.message.message_id).gid) + uploaded_bytes = 0 + while response is None: + status, response = drive_file.next_chunk() + time_lapsed = time.time() - self.start_time + + if status: + # The iconic formula of speed = distance / time :) + LOGGER.info(status.progress() * 100) + chunk_size = status.total_size*status.progress() - uploaded_bytes + uploaded_bytes = status.total_size*status.progress() + download_dict[self.__listener.uid].uploaded_bytes += chunk_size + download_dict[self.__listener.uid].upload_time = time_lapsed + self.__listener.onUploadProgress(_list, index) # Insert new permissions - self.__service.permissions().create(fileId=drive_file['id'], body=permissions).execute() + self.__service.permissions().create(fileId=response['id'], body=permissions).execute() # Define file instance and get url for download - drive_file = self.__service.files().get(fileId=drive_file['id']).execute() + drive_file = self.__service.files().get(fileId=response['id']).execute() download_url = self.__G_DRIVE_BASE_DOWNLOAD_URL.format(drive_file.get('id')) return download_url @@ -61,17 +79,20 @@ class GoogleDriveHelper: self.__listener.onUploadStarted(_list, index) file_dir = f"{DOWNLOAD_DIR}{self.__listener.message.message_id}" file_path = f"{file_dir}/{file_name}" - link = None LOGGER.info("Uploading File: " + file_name) + self.start_time = time.time() if os.path.isfile(file_path): - mime_type = get_mime_type(file_path) try: + mime_type = get_mime_type(file_path) g_drive_link = self.upload_file(file_path, file_name, mime_type, parent_id) LOGGER.info("Uploaded To G-Drive: " + file_path) link = g_drive_link except Exception as e: LOGGER.error(str(e)) - self.__listener.onUploadError(str(e), _list, index) + e_str = str(e).replace('<', '') + e_str = e_str.replace('>', '') + self.__listener.onUploadError(e_str, _list, index) + return else: try: dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_id) @@ -80,7 +101,10 @@ class GoogleDriveHelper: link = f"https://drive.google.com/folderview?id={dir_id}" except Exception as e: LOGGER.error(str(e)) - self.__listener.onUploadError(str(e), _list, index) + e_str = str(e).replace('<', '') + e_str = e_str.replace('>', '') + self.__listener.onUploadError(e_str, _list, index) + return LOGGER.info(download_dict) self.__listener.onUploadComplete(link, _list, index) LOGGER.info("Deleting downloaded file/folder..") diff --git a/bot/helper/listeners.py b/bot/helper/listeners.py index d19fe16..5840ed0 100644 --- a/bot/helper/listeners.py +++ b/bot/helper/listeners.py @@ -3,6 +3,7 @@ class MirrorListeners: self.context = context self.update = update self.message = update.message + self.uid = self.message.message_id self.reply_message = reply_message def onDownloadStarted(self, link: str): @@ -20,6 +21,9 @@ class MirrorListeners: def onUploadStarted(self, progress_status_list: list, index: int): raise NotImplementedError + def onUploadProgress(self, progress: list, index: int): + raise NotImplementedError + def onUploadComplete(self, link: str, progress_status_list: list, index: int): raise NotImplementedError diff --git a/bot/mirror.py b/bot/mirror.py index 26e9cc6..c1b6ef4 100644 --- a/bot/mirror.py +++ b/bot/mirror.py @@ -5,8 +5,7 @@ from bot import LOGGER, dispatcher from bot.helper import fs_utils from bot import download_dict, status_reply_dict from bot.helper.message_utils import * -from bot.helper.bot_utils import get_readable_message, KillThreadException -from bot.helper.download_status import DownloadStatus +from bot.helper.bot_utils import get_readable_message, KillThreadException, MirrorStatus class MirrorListener(listeners.MirrorListeners): @@ -17,7 +16,7 @@ class MirrorListener(listeners.MirrorListeners): LOGGER.info("Adding link: " + link) def onDownloadProgress(self, progress_status_list: list, index: int): - if progress_status_list[index].status() == DownloadStatus.STATUS_CANCELLED: + if progress_status_list[index].status() == MirrorStatus.STATUS_CANCELLED: raise KillThreadException('Mirror cancelled by user') msg = get_readable_message(progress_status_list) # LOGGER.info("Editing message") @@ -65,6 +64,10 @@ class MirrorListener(listeners.MirrorListeners): del download_dict[self.message.message_id] fs_utils.clean_download(progress_status[index].path()) + def onUploadProgress(self, progress: list, index: int): + msg = get_readable_message(progress) + editMessage(msg, self.context, self.reply_message) + @run_async def mirror(update, context):