From 54d1b96aa2be97c2e5d27485a71e1df04e1007ef Mon Sep 17 00:00:00 2001 From: anasty17 Date: Tue, 12 Dec 2023 00:40:15 +0200 Subject: [PATCH] Minor changes Signed-off-by: anasty17 --- bot/helper/common.py | 22 ++++++++-------- bot/helper/ext_utils/files_utils.py | 7 +++-- bot/helper/listeners/task_listener.py | 4 +-- bot/helper/mirror_utils/gdrive_utils/clone.py | 16 +++++------- .../mirror_utils/gdrive_utils/download.py | 12 +++------ .../mirror_utils/gdrive_utils/helper.py | 3 +-- bot/helper/mirror_utils/gdrive_utils/list.py | 4 +-- .../mirror_utils/gdrive_utils/upload.py | 26 +++++++------------ .../mirror_utils/rclone_utils/transfer.py | 20 +++++--------- bot/helper/mirror_utils/telegram_uploader.py | 10 ++----- bot/modules/clone.py | 9 +++---- bot/modules/mirror_leech.py | 7 ++--- bot/modules/ytdlp.py | 5 ++-- 13 files changed, 58 insertions(+), 87 deletions(-) diff --git a/bot/helper/common.py b/bot/helper/common.py index ae1f8dbf066..9de936c0ad2 100644 --- a/bot/helper/common.py +++ b/bot/helper/common.py @@ -56,32 +56,28 @@ class TaskConfig: - def __init__(self, message): - self.message = message + def __init__(self): self.mid = self.message.id self.user = self.message.from_user or self.message.sender_chat self.user_id = self.user.id self.user_dict = user_data.get(self.user_id, {}) - self.sameDir = {} - self.bulk = [] self.dir = f"{DOWNLOAD_DIR}{self.mid}" self.link = "" self.upDest = "" self.rcFlags = "" - self.options = "" self.tag = "" self.name = "" self.session = "" self.newDir = "" - self.multiTag = 0 self.splitSize = 0 self.maxSplitSize = 0 self.multi = 0 + self.isLeech = False + self.isQbit = False + self.isClone = False + self.isYtDlp = False self.equalSplits = False self.userTransmission = False - self.isClone = False - self.isQbit = False - self.isLeech = False self.extract = False self.compress = False self.select = False @@ -89,15 +85,14 @@ def __init__(self, message): self.compress = False self.extract = False self.join = False - self.isYtDlp = False self.privateLink = False self.stopDuplicate = False self.sampleVideo = False self.screenShots = False self.as_doc = False self.suproc = None - self.client = None self.thumb = None + self.extension_filter = [] self.isSuperChat = self.message.chat.type.name in ["SUPERGROUP", "CHANNEL"] def getTokenPath(self, dest): @@ -137,6 +132,11 @@ async def isTokenExists(self, path, status): raise ValueError(f"SAccounts or token.pickle: {token_path} not Exists!") async def beforeStart(self): + self.extension_filter = ( + self.user_dict.get("excluded_extensions") or GLOBAL_EXTENSION_FILTER + if "excluded_extensions" not in self.user_dict + else ["aria2", "!qB"] + ) if not self.isYtDlp: if self.link not in ["rcl", "gdl"]: await self.isTokenExists(self.link, "dl") diff --git a/bot/helper/ext_utils/files_utils.py b/bot/helper/ext_utils/files_utils.py index bc3443ffd05..8745972da20 100644 --- a/bot/helper/ext_utils/files_utils.py +++ b/bot/helper/ext_utils/files_utils.py @@ -1,14 +1,13 @@ from os import walk, path as ospath from aiofiles.os import remove as aioremove, path as aiopath, listdir, rmdir, makedirs from aioshutil import rmtree as aiormtree -from shutil import rmtree from magic import Magic from re import split as re_split, I, search as re_search from subprocess import run as srun from sys import exit as sexit from .exceptions import NotSupportedExtractionArchive -from bot import aria2, LOGGER, DOWNLOAD_DIR, get_client, GLOBAL_EXTENSION_FILTER +from bot import aria2, LOGGER, DOWNLOAD_DIR, get_client from bot.helper.ext_utils.bot_utils import sync_to_async, async_to_sync, cmd_exec ARCH_EXT = [ @@ -138,13 +137,13 @@ async def get_path_size(path): return total_size -async def count_files_and_folders(path): +async def count_files_and_folders(path, extension_filter): total_files = 0 total_folders = 0 for _, dirs, files in await sync_to_async(walk, path): total_files += len(files) for f in files: - if f.endswith(tuple(GLOBAL_EXTENSION_FILTER)): + if f.endswith(tuple(extension_filter)): total_files -= 1 total_folders += len(dirs) return total_folders, total_files diff --git a/bot/helper/listeners/task_listener.py b/bot/helper/listeners/task_listener.py index 356fbee92fd..e5a16011561 100644 --- a/bot/helper/listeners/task_listener.py +++ b/bot/helper/listeners/task_listener.py @@ -47,8 +47,8 @@ class TaskListener(TaskConfig): - def __init__(self, message): - super().__init__(message) + def __init__(self): + super().__init__() async def clean(self): try: diff --git a/bot/helper/mirror_utils/gdrive_utils/clone.py b/bot/helper/mirror_utils/gdrive_utils/clone.py index 868e0d5241c..e7a7fe38e66 100644 --- a/bot/helper/mirror_utils/gdrive_utils/clone.py +++ b/bot/helper/mirror_utils/gdrive_utils/clone.py @@ -10,7 +10,6 @@ RetryError, ) -from bot import GLOBAL_EXTENSION_FILTER from bot.helper.ext_utils.bot_utils import async_to_sync from bot.helper.mirror_utils.gdrive_utils.helper import GoogleDriveHelper @@ -19,8 +18,9 @@ class gdClone(GoogleDriveHelper): def __init__(self, listener): - super().__init__(listener) + self.listener = listener self._start_time = time() + super().__init__() self.is_cloning = True self.user_setting() @@ -109,19 +109,17 @@ def _cloneFolder(self, folder_name, folder_id, dest_id): files = self.getFilesByFolderId(folder_id) if len(files) == 0: return dest_id - if self.listener.user_dict.get("excluded_extensions", False): - extension_filter = self.listener.user_dict["excluded_extensions"] - elif "excluded_extensions" not in self.listener.user_dict: - extension_filter = GLOBAL_EXTENSION_FILTER - else: - extension_filter = ["aria2", "!qB"] for file in files: if file.get("mimeType") == self.G_DRIVE_DIR_MIME_TYPE: self.total_folders += 1 file_path = ospath.join(folder_name, file.get("name")) current_dir_id = self.create_directory(file.get("name"), dest_id) self._cloneFolder(file_path, file.get("id"), current_dir_id) - elif not file.get("name").lower().endswith(tuple(extension_filter)): + elif ( + not file.get("name") + .lower() + .endswith(tuple(self.listener.extension_filter)) + ): self.total_files += 1 self._copyFile(file.get("id"), dest_id) self.proc_bytes += int(file.get("size", 0)) diff --git a/bot/helper/mirror_utils/gdrive_utils/download.py b/bot/helper/mirror_utils/gdrive_utils/download.py index 59e1af3b4ab..42ca9468b4b 100644 --- a/bot/helper/mirror_utils/gdrive_utils/download.py +++ b/bot/helper/mirror_utils/gdrive_utils/download.py @@ -11,7 +11,6 @@ RetryError, ) -from bot import GLOBAL_EXTENSION_FILTER from bot.helper.ext_utils.bot_utils import setInterval from bot.helper.ext_utils.bot_utils import async_to_sync from bot.helper.mirror_utils.gdrive_utils.helper import GoogleDriveHelper @@ -21,9 +20,10 @@ class gdDownload(GoogleDriveHelper): def __init__(self, listener, path): - super().__init__(listener) + self.listener = listener self._updater = None self._path = path + super().__init__() self.is_downloading = True def download(self): @@ -70,12 +70,6 @@ def _download_folder(self, folder_id, path, folder_name): result = self.getFilesByFolderId(folder_id) if len(result) == 0: return - if self.listener.user_dict.get("excluded_extensions", False): - extension_filter = self.listener.user_dict["excluded_extensions"] - elif "excluded_extensions" not in self.listener.user_dict: - extension_filter = GLOBAL_EXTENSION_FILTER - else: - extension_filter = ["aria2", "!qB"] result = sorted(result, key=lambda k: k["name"]) for item in result: file_id = item["id"] @@ -90,7 +84,7 @@ def _download_folder(self, folder_id, path, folder_name): self._download_folder(file_id, path, filename) elif not ospath.isfile( f"{path}{filename}" - ) and not filename.lower().endswith(tuple(extension_filter)): + ) and not filename.lower().endswith(tuple(self.listener.extension_filter)): self._download_file(file_id, path, filename, mime_type) if self.is_cancelled: break diff --git a/bot/helper/mirror_utils/gdrive_utils/helper.py b/bot/helper/mirror_utils/gdrive_utils/helper.py index 34eaa6debb4..d4b5a11775f 100644 --- a/bot/helper/mirror_utils/gdrive_utils/helper.py +++ b/bot/helper/mirror_utils/gdrive_utils/helper.py @@ -21,7 +21,7 @@ class GoogleDriveHelper: - def __init__(self, listener=None): + def __init__(self): self._OAUTH_SCOPE = ["https://www.googleapis.com/auth/drive"] self.token_path = "token.pickle" self.G_DRIVE_DIR_MIME_TYPE = "application/vnd.google-apps.folder" @@ -37,7 +37,6 @@ def __init__(self, listener=None): self.sa_count = 1 self.sa_number = 100 self.alt_auth = False - self.listener = listener self.service = None self.total_files = 0 self.total_folders = 0 diff --git a/bot/helper/mirror_utils/gdrive_utils/list.py b/bot/helper/mirror_utils/gdrive_utils/list.py index 88b0797e85c..0ff7ae02b3e 100644 --- a/bot/helper/mirror_utils/gdrive_utils/list.py +++ b/bot/helper/mirror_utils/gdrive_utils/list.py @@ -113,7 +113,7 @@ async def id_updates(_, query, obj): class gdriveList(GoogleDriveHelper): def __init__(self, listener): - super().__init__(listener) + self.listener = listener self._token_user = False self._token_owner = False self._sa_owner = False @@ -121,7 +121,6 @@ def __init__(self, listener): self._time = time() self._timeout = 240 self.drives = [] - self.is_cancelled = False self.query_proc = False self.item_type = "folders" self.event = Event() @@ -132,6 +131,7 @@ def __init__(self, listener): self.items_list = [] self.iter_start = 0 self.page_step = 1 + super().__init__() @new_thread async def _event_handler(self): diff --git a/bot/helper/mirror_utils/gdrive_utils/upload.py b/bot/helper/mirror_utils/gdrive_utils/upload.py index 3127b6cc437..e12020bdce7 100644 --- a/bot/helper/mirror_utils/gdrive_utils/upload.py +++ b/bot/helper/mirror_utils/gdrive_utils/upload.py @@ -10,7 +10,7 @@ RetryError, ) -from bot import config_dict, GLOBAL_EXTENSION_FILTER +from bot import config_dict from bot.helper.ext_utils.files_utils import get_mime_type from bot.helper.ext_utils.bot_utils import async_to_sync, setInterval from bot.helper.mirror_utils.gdrive_utils.helper import GoogleDriveHelper @@ -20,10 +20,11 @@ class gdUpload(GoogleDriveHelper): def __init__(self, listener, path): - super().__init__(listener) + self.listener = listener self._updater = None self._path = path self._is_errored = False + super().__init__() self.is_uploading = True def user_setting(self): @@ -44,14 +45,8 @@ def upload(self, size): LOGGER.info(f"Uploading: {self._path}") self._updater = setInterval(self.update_interval, self.progress) try: - if self.listener.user_dict.get("excluded_extensions", False): - extension_filter = self.listener.user_dict["excluded_extensions"] - elif "excluded_extensions" not in self.listener.user_dict: - extension_filter = GLOBAL_EXTENSION_FILTER - else: - extension_filter = ["aria2", "!qB"] if ospath.isfile(self._path): - if self._path.lower().endswith(tuple(extension_filter)): + if self._path.lower().endswith(tuple(self.listener.extension_filter)): raise Exception( "This file extension is excluded by extension filter!" ) @@ -71,9 +66,10 @@ def upload(self, size): else: mime_type = "Folder" dir_id = self.create_directory( - ospath.basename(ospath.abspath(self.listener.name)), self.listener.upDest + ospath.basename(ospath.abspath(self.listener.name)), + self.listener.upDest, ) - result = self._upload_dir(self._path, dir_id, extension_filter) + result = self._upload_dir(self._path, dir_id) if result is None: raise Exception("Upload has been manually cancelled!") link = self.G_DRIVE_DIR_BASE_DOWNLOAD_URL.format(dir_id) @@ -108,7 +104,7 @@ def upload(self, size): dir_id=self.getIdFromUrl(link), ) - def _upload_dir(self, input_directory, dest_id, extension_filter): + def _upload_dir(self, input_directory, dest_id): list_dirs = listdir(input_directory) if len(list_dirs) == 0: return dest_id @@ -117,11 +113,9 @@ def _upload_dir(self, input_directory, dest_id, extension_filter): current_file_name = ospath.join(input_directory, item) if ospath.isdir(current_file_name): current_dir_id = self.create_directory(item, dest_id) - new_id = self._upload_dir( - current_file_name, current_dir_id, extension_filter - ) + new_id = self._upload_dir(current_file_name, current_dir_id) self.total_folders += 1 - elif not item.lower().endswith(tuple(extension_filter)): + elif not item.lower().endswith(tuple(self.listener.extension_filter)): mime_type = get_mime_type(current_file_name) file_name = current_file_name.split("/")[-1] self._upload_file(current_file_name, file_name, mime_type, dest_id) diff --git a/bot/helper/mirror_utils/rclone_utils/transfer.py b/bot/helper/mirror_utils/rclone_utils/transfer.py index 275e4e5fcb0..6200150e9ce 100644 --- a/bot/helper/mirror_utils/rclone_utils/transfer.py +++ b/bot/helper/mirror_utils/rclone_utils/transfer.py @@ -8,7 +8,7 @@ from random import randrange from logging import getLogger -from bot import config_dict, GLOBAL_EXTENSION_FILTER +from bot import config_dict from bot.helper.ext_utils.bot_utils import cmd_exec, sync_to_async from bot.helper.ext_utils.files_utils import get_mime_type, count_files_and_folders @@ -31,8 +31,6 @@ def __init__(self, listener): self._sa_count = 1 self._sa_index = 0 self._sa_number = 0 - self.extension_filter = ["aria2", "!qB"] - self.user_settings() @property def transferred_size(self): @@ -54,14 +52,6 @@ def eta(self): def size(self): return self._size - def user_settings(self): - if self._listener.user_dict.get("excluded_extensions", False): - self.extension_filter = self._listener.user_dict["excluded_extensions"] - elif "excluded_extensions" not in self._listener.user_dict: - self.extension_filter = GLOBAL_EXTENSION_FILTER - else: - self.extension_filter = ["aria2", "!qB"] - async def _progress(self): while not (self._proc is None or self._is_cancelled): try: @@ -278,10 +268,12 @@ async def upload(self, path, size): if await aiopath.isdir(path): mime_type = "Folder" - folders, files = await count_files_and_folders(path) + folders, files = await count_files_and_folders( + path, self._listener.extension_filter + ) rc_path += f"/{self._listener.name}" if rc_path else self._listener.name else: - if path.lower().endswith(tuple(self.extension_filter)): + if path.lower().endswith(tuple(self._listener.extension_filter)): await self._listener.onUploadError( "This file extension is excluded by extension filter!" ) @@ -427,7 +419,7 @@ async def clone(self, config_path, src_remote, src_path, mime_type): return None, None def _getUpdatedCommand(self, config_path, source, destination, method): - ext = "*.{" + ",".join(self.extension_filter) + "}" + ext = "*.{" + ",".join(self._listener.extension_filter) + "}" cmd = f'rclone {method} --fast-list --config {config_path} -P "{source}" "{destination}" --exclude "{ext}"' cmd += " --ignore-case --low-level-retries 1 -M --log-file rlog.txt --log-level DEBUG" if rcflags := self._listener.rcFlags or config_dict["RCLONE_FLAGS"]: diff --git a/bot/helper/mirror_utils/telegram_uploader.py b/bot/helper/mirror_utils/telegram_uploader.py index 1f04208a724..1fe8ab75cfc 100644 --- a/bot/helper/mirror_utils/telegram_uploader.py +++ b/bot/helper/mirror_utils/telegram_uploader.py @@ -22,7 +22,7 @@ from natsort import natsorted from aioshutil import copy -from bot import config_dict, GLOBAL_EXTENSION_FILTER, user +from bot import config_dict, user from bot.helper.ext_utils.files_utils import clean_unwanted, is_archive, get_base_name from bot.helper.ext_utils.bot_utils import sync_to_async from bot.helper.ext_utils.media_utils import ( @@ -223,18 +223,12 @@ async def upload(self, o_files, m_size, size): res = await self._msg_to_reply() if not res: return - if self._listener.user_dict.get("excluded_extensions", False): - extension_filter = self._listener.user_dict["excluded_extensions"] - elif "excluded_extensions" not in self._listener.user_dict: - extension_filter = GLOBAL_EXTENSION_FILTER - else: - extension_filter = ["aria2", "!qB"] for dirpath, _, files in sorted(await sync_to_async(walk, self._path)): if dirpath.endswith("/yt-dlp-thumb"): continue for file_ in natsorted(files): self._up_path = ospath.join(dirpath, file_) - if file_.lower().endswith(tuple(extension_filter)): + if file_.lower().endswith(tuple(self._listener.extension_filter)): if not self._listener.seed or self._listener.newDir: await aioremove(self._up_path) continue diff --git a/bot/modules/clone.py b/bot/modules/clone.py index ede7dd727f1..cf84feffe8f 100644 --- a/bot/modules/clone.py +++ b/bot/modules/clone.py @@ -46,21 +46,20 @@ def __init__( message, _=None, __=None, - sameDir=None, + ___=None, bulk=None, multiTag=None, options="", ): - if sameDir is None: - sameDir = {} if bulk is None: bulk = [] - super().__init__(message) + self.message = message self.client = client self.multiTag = multiTag self.options = options - self.sameDir = sameDir + self.sameDir = {} self.bulk = bulk + super().__init__() self.isClone = True @new_task diff --git a/bot/modules/mirror_leech.py b/bot/modules/mirror_leech.py index 6a7d5f45d11..eeda08d320f 100644 --- a/bot/modules/mirror_leech.py +++ b/bot/modules/mirror_leech.py @@ -56,14 +56,15 @@ def __init__( sameDir = {} if bulk is None: bulk = [] - super().__init__(message) + self.message = message self.client = client - self.isQbit = isQbit - self.isLeech = isLeech self.multiTag = multiTag self.options = options self.sameDir = sameDir self.bulk = bulk + super().__init__() + self.isQbit = isQbit + self.isLeech = isLeech @new_task async def newEvent(self): diff --git a/bot/modules/ytdlp.py b/bot/modules/ytdlp.py index 2a4baab2ec5..30ad18456f4 100644 --- a/bot/modules/ytdlp.py +++ b/bot/modules/ytdlp.py @@ -273,14 +273,15 @@ def __init__( sameDir = {} if bulk is None: bulk = [] - super().__init__(message) + self.message = message self.client = client self.isLeech = isLeech - self.isYtDlp = True self.multiTag = multiTag self.options = options self.sameDir = sameDir self.bulk = bulk + super().__init__() + self.isYtDlp = True @new_task async def newEvent(self):