Introduce DownloadStatus class for storing stats of downloads

Signed-off-by: lzzy12 <jhashivam2020@gmail.com>

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

bot/helper/download_status.py

@@ -1,8 +1,54 @@
-class Message:
+from bot import aria2
-    pass
 
-class Download:
+
-    def __init__(self, gid, download):
+def get_download(gid):
-        self.isDownloading = download.is_complete
+    return aria2.get_download(gid)
-        self.gid = gid
+
-        self.download = download
+
+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

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("<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))

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

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

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 += "<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"
+        msg = get_readable_message(progress_status_list)
-        for progress_str in progress_str_list:
-            msg += progress_str + '\n\n'
         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):
+    def onDownloadComplete(self, progress_status_list, index: int):
-        msg = "Downloaded"
+        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):
+    def onUploadStarted(self, progress_status_list: list, index: int):
-        msg = "<i>" + filename + "</i>:- Uploading."
+        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):
+    def onUploadComplete(self, link: str, progress_status_list: list, index: int):
-        msg = '<a href="{}">{}</a>'.format(link, file_name)
+        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,