Introduce DownloadStatus class for storing stats of downloads

Signed-off-by: lzzy12 <jhashivam2020@gmail.com>
This commit is contained in:
lzzy12 2019-09-26 22:34:37 +05:30
parent f0efb4357d
commit 84482969f1
6 changed files with 128 additions and 56 deletions

24
bot/helper/bot_utils.py Normal file
View File

@ -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

View File

@ -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

View File

@ -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("<b>" + download.name + "</b>:- "
+ 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))

View File

@ -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

View File

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

View File

@ -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 += "<b>Name:</b> {}\n" \
"<b>status:</b> {}\n" \
"<b>Downloaded:</b> {} of {}\n" \
"<b>Speed:</b> {}\n" \
"<b>ETA:</b> {}\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 = "<i>" + filename + "</i>:- 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 = '<a href="{}">{}</a>'.format(link, file_name)
def onUploadComplete(self, link: str, progress_status_list: list, index: int):
msg = '<a href="{}">{}</a>'.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,