evanferrao
/
dump-mirrorbot
mirror of https://github.com/Box-boi/dump-mirrorbot.git
1
0
Fork 0
Code Issues Actions Packages Projects Releases Wiki Activity

Improvements in cancelling mirrors 

Signed-off-by: lzzy12 <jhashivam2020@gmail.com>
This commit is contained in:
lzzy12 2020-04-04 22:58:11 +05:30
parent 708e6df87a
commit b58eb87b11
2 changed files with 30 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
bot/helper/mirror_utils/download_utils/aria2_download.py
View File

@ -12,8 +12,8 @@ class AriaDownloadHelper(DownloadHelper):
def __init__(self, listener):
super().__init__()
self.gid = None
self._listener = listener
self._resource_lock = threading.Lock()
self.__listener = listener
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].is_torrent =True
download_dict[self.__listener.uid] = AriaDownloadStatus(self.gid, self.__listener)
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])

36
bot/modules/cancel_mirror.py
View File

@ -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":
aria2.pause([dlDetails.download()])
if not dlDetails.status() == MirrorStatus.STATUS_UPLOADING\
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)