Improvements in cancelling mirrors

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

bot/helper/mirror_utils/download_utils/aria2_download.py

@@ -12,8 +12,8 @@ class AriaDownloadHelper(DownloadHelper):
     def __init__(self, listener):
         super().__init__()
         self.gid = None
-        self._listener = listener
+        self.__listener = listener
-        self._resource_lock = threading.Lock()
+        self._resource_lock = threading.RLock()
 
     def __onDownloadStarted(self, api, gid):
         with self._resource_lock:
@@ -28,22 +28,22 @@ class AriaDownloadHelper(DownloadHelper):
                 if api.get_download(gid).followed_by_ids:
                     self.gid = api.get_download(gid).followed_by_ids[0]
                     with download_dict_lock:
-                        download_dict[self._listener.uid] = AriaDownloadStatus(self.gid, self._listener)
+                        download_dict[self.__listener.uid] = AriaDownloadStatus(self.gid, self.__listener)
-                        download_dict[self._listener.uid].is_torrent =True
+                        download_dict[self.__listener.uid].is_torrent =True
                     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:
             LOGGER.info("Called onDownloadPause")
-            self._listener.onDownloadError('Download stopped by user!')
+            self.__listener.onDownloadError('Download stopped by user!')
 
     def __onDownloadStopped(self, api, gid):
         if self.gid == gid:
             LOGGER.info("Called on_download_stop")
-            self._listener.onDownloadError('Download stopped by user!')
+            self.__listener.onDownloadError('Download stopped by user!')
 
     def __onDownloadError(self, api, gid):
         with self._resource_lock:
@@ -51,7 +51,7 @@ class AriaDownloadHelper(DownloadHelper):
                 download = api.get_download(gid)
                 error = download.error_message
                 LOGGER.info(f"Download Error: {error}")
-                self._listener.onDownloadError(error)
+                self.__listener.onDownloadError(error)
 
     def add_download(self, link: str, path):
         if is_magnet(link):
@@ -60,9 +60,9 @@ class AriaDownloadHelper(DownloadHelper):
             download = aria2.add_uris([link], {'dir': path})
         self.gid = download.gid
         with download_dict_lock:
-            download_dict[self._listener.uid] = AriaDownloadStatus(self.gid, self._listener)
+            download_dict[self.__listener.uid] = AriaDownloadStatus(self.gid, self.__listener)
         if download.error_message:
-            self._listener.onDownloadError(download.error_message)
+            self.__listener.onDownloadError(download.error_message)
             return
         LOGGER.info(f"Started: {self.gid} DIR:{download.dir} ")
         aria2.listen_to_notifications(threaded=True, on_download_start=self.__onDownloadStarted,
@@ -70,3 +70,14 @@ class AriaDownloadHelper(DownloadHelper):
                                       on_download_pause=self.__onDownloadPause,
                                       on_download_stop=self.__onDownloadStopped,
                                       on_download_complete=self.__onDownloadComplete)
+
+    def cancel_download(self):
+        download = aria2.get_download(self.gid)
+        if download.is_queued:
+            aria2.remove([download])
+            self.__listener.onDownloadError("Cancelled by user")
+            return
+        if len(download.followed_by_ids) != 0:
+            downloads = aria2.get_downloads(download.followed_by_ids)
+            aria2.pause(downloads)
+        aria2.pause([download])

bot/modules/cancel_mirror.py

@@ -1,22 +1,18 @@
-from time import sleep
-
 from telegram.ext import CommandHandler, run_async
 
-from bot import download_dict, aria2, dispatcher, download_dict_lock, DOWNLOAD_DIR
+from bot import download_dict, dispatcher, download_dict_lock, DOWNLOAD_DIR
 from bot.helper.ext_utils.fs_utils import clean_download
 from bot.helper.telegram_helper.bot_commands import BotCommands
 from bot.helper.telegram_helper.filters import CustomFilters
 from bot.helper.telegram_helper.message_utils import *
-from ..helper.mirror_utils.download_utils.telegram_downloader import TelegramDownloadHelper
 
 from time import sleep
 from bot.helper.ext_utils.bot_utils import getDownloadByGid
-from ..helper.mirror_utils.download_utils.youtube_dl_download_helper import YoutubeDLHelper
+from ..helper.ext_utils.bot_utils import MirrorStatus
 
 
 @run_async
 def cancel_mirror(bot, update):
-    # TODO: Rewrite this for good
     args = update.message.text.split(" ", maxsplit=1)
     mirror_message = None
     if len(args) > 1:
@@ -27,7 +23,7 @@ def cancel_mirror(bot, update):
             return
         with download_dict_lock:
             keys = list(download_dict.keys())
-        mirror_message = dl._listener.message
+        mirror_message = dl.__listener.message
     elif update.message.reply_to_message:
         mirror_message = update.message.reply_to_message
         with download_dict_lock:
@@ -45,25 +41,14 @@ def cancel_mirror(bot, update):
                       "used to start the download or /cancel gid to cancel it!"
                 sendMessage(msg, bot, update)
                 return
-    if dl.status() == "Uploading":
+    if dl.status() == MirrorStatus.STATUS_UPLOADING:
         sendMessage("Upload in Progress, Don't Cancel it.", bot, update)
         return
-    elif dl.status() == "Archiving":
+    elif dl.status() == MirrorStatus.STATUS_ARCHIVING:
         sendMessage("Archival in Progress, Don't Cancel it.", bot, update)
         return
-    elif dl.status() != "Queued":
-        download = dl.download()
-        if isinstance(download, TelegramDownloadHelper):
-            download.cancel_download()
-        if isinstance(download, YoutubeDLHelper):
-            download.cancel_download()
-        else:
-            if len(download.followed_by_ids) != 0:
-                downloads = aria2.get_downloads(download.followed_by_ids)
-                aria2.pause(downloads)
-            aria2.pause([download])
     else:
-        dl._listener.onDownloadError("Download stopped by user!")
+        dl.download().cancel_download()
     sleep(1)  # Wait a Second For Aria2 To free Resources.
     clean_download(f'{DOWNLOAD_DIR}{mirror_message.message_id}/')
 
@@ -73,13 +58,10 @@ def cancel_all(update, bot):
     with download_dict_lock:
         count = 0
         for dlDetails in list(download_dict.values()):
-            if not dlDetails.status() == "Uploading" or dlDetails.status() == "Archiving":
+            if not dlDetails.status() == MirrorStatus.STATUS_UPLOADING\
-                aria2.pause([dlDetails.download()])
+                    or not dlDetails.status() == MirrorStatus.STATUS_ARCHIVING:
+                dlDetails.cancel_download()
                 count += 1
-                continue
-            if dlDetails.status() == "Queued":
-                count += 1
-                dlDetails._listener.onDownloadError("Download Manually Cancelled By user.")
     delete_all_messages()
     sendMessage(f'Cancelled {count} downloads!', update, bot)