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):