diff --git a/bot/helper/drive_utils/gdriveTools.py b/bot/helper/drive_utils/gdriveTools.py index 4789284d..1ac83a78 100644 --- a/bot/helper/drive_utils/gdriveTools.py +++ b/bot/helper/drive_utils/gdriveTools.py @@ -340,18 +340,19 @@ def __copyFile(self, file_id, dest_id): except HttpError as err: if err.resp.get('content-type', '').startswith('application/json'): reason = json.loads(err.content).get('error').get('errors')[0].get('reason') - if reason in ['userRateLimitExceeded', 'dailyLimitExceeded']: - if USE_SERVICE_ACCOUNTS: - if self.__sa_count == SERVICE_ACCOUNTS_NUMBER: - LOGGER.info("SA switch limit exceeded") - raise err - else: - self.__switchServiceAccount() - return self.__copyFile(file_id, dest_id) - else: - LOGGER.error(f"Warning: {reason}") + if reason not in ['userRateLimitExceeded', 'dailyLimitExceeded']: + raise err + if USE_SERVICE_ACCOUNTS: + if self.__sa_count >= SERVICE_ACCOUNTS_NUMBER: + LOGGER.info("SA switch limit exceeded") raise err + else: + if self.__is_cancelled: + return + self.__switchServiceAccount() + return self.__copyFile(file_id, dest_id) else: + LOGGER.error(f"Warning: {reason}") raise err def __cloneFolder(self, name, local_path, folder_id, dest_id): @@ -529,8 +530,15 @@ def __upload_file(self, file_path, file_name, mime_type, dest_id): if reason not in ['userRateLimitExceeded', 'dailyLimitExceeded']: raise err if USE_SERVICE_ACCOUNTS: - self.__switchServiceAccount() - return self.__upload_file(file_path, file_name, mime_type, dest_id) + if self.__sa_count >= SERVICE_ACCOUNTS_NUMBER: + LOGGER.info("SA switch limit exceeded") + raise err + else: + if self.__is_cancelled: + return + self.__switchServiceAccount() + LOGGER.info(f"Warning: {reason}") + return self.__upload_file(file_path, file_name, mime_type, dest_id) else: LOGGER.error(f"Warning: {reason}") raise err @@ -637,11 +645,14 @@ def __download_file(self, file_id, path, filename, mime_type): if reason not in ['downloadQuotaExceeded', 'dailyLimitExceeded']: raise err if USE_SERVICE_ACCOUNTS: - if self.__sa_count == SERVICE_ACCOUNTS_NUMBER: + if self.__sa_count >= SERVICE_ACCOUNTS_NUMBER: LOGGER.info("SA switch limit exceeded") raise err else: + if self.__is_cancelled: + return self.__switchServiceAccount() + LOGGER.info(f"Warning: {reason}") return self.__download_file(file_id, path, filename, mime_type) else: LOGGER.error(f"Warning: {reason}") @@ -682,7 +693,7 @@ def download(self, link): if meta.get("mimeType") == self.__G_DRIVE_DIR_MIME_TYPE: self.__download_folder(file_id, self.__path, self.name) else: - os.makedirs(self.__path) + os.makedirs(self.__path, exist_ok=True) self.__download_file(file_id, self.__path, self.name, meta.get('mimeType')) except Exception as err: if isinstance(err, RetryError): diff --git a/bot/helper/ext_utils/fs_utils.py b/bot/helper/ext_utils/fs_utils.py index e7aef7fb..ec7a9eae 100644 --- a/bot/helper/ext_utils/fs_utils.py +++ b/bot/helper/ext_utils/fs_utils.py @@ -66,10 +66,9 @@ def get_path_size(path: str): return total_size def get_base_name(orig_path: str): - ext = [ext for ext in ARCH_EXT if orig_path.lower().endswith(ext)] - if ext: + if ext:= [ext for ext in ARCH_EXT if orig_path.lower().endswith(ext)]: ext = ext[0] - return re.split(ext + '$', orig_path, maxsplit=1, flags=re.I)[0] + return re.split(f'{ext}$', orig_path, maxsplit=1, flags=re.I)[0] else: raise ArchiveExceptionHandler("Unsupported archive file") diff --git a/bot/modules/archive.py b/bot/modules/archive.py index 46b2c994..8dc89402 100644 --- a/bot/modules/archive.py +++ b/bot/modules/archive.py @@ -1,7 +1,6 @@ import os import re import requests -import shutil import subprocess import threading @@ -49,7 +48,7 @@ def onDownloadComplete(self): gid = download.gid() if name == 'None' or not os.path.exists(f'{self.dir}/{name}'): name = os.listdir(self.dir)[-1] - m_path = f'{self.dir}/{name}' + m_path = f"{self.dir}/{name}" size = get_path_size(m_path) if self.is_compress: path = f"{m_path}.zip" @@ -145,10 +144,8 @@ def onUploadComplete(self, link: str, size, files, folders, typ, name): sendMessage(msg, self.bot, self.message) clean_download(self.dir) with download_dict_lock: - try: + if self.uid in download_dict.keys(): del download_dict[self.uid] - except Exception as err: - LOGGER.error(err) count = len(download_dict) if count == 0: self.clean() @@ -159,10 +156,8 @@ def onDownloadError(self, error): error = error.replace('<', '').replace('>', '') clean_download(self.dir) with download_dict_lock: - try: + if self.uid in download_dict.keys(): del download_dict[self.uid] - except Exception as err: - LOGGER.error(err) count = len(download_dict) sendMessage(error, self.bot, self.message) if count == 0: @@ -174,10 +169,8 @@ def onUploadError(self, error): error = error.replace('<', '').replace('>', '') clean_download(self.dir) with download_dict_lock: - try: + if self.uid in download_dict.keys(): del download_dict[self.uid] - except Exception as err: - LOGGER.error(err) count = len(download_dict) sendMessage(error, self.bot, self.message) if count == 0: diff --git a/requirements.txt b/requirements.txt index 96d431f9..1d5ac517 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,7 +5,7 @@ psutil pymongo python-dotenv python-magic -python-telegram-bot +python-telegram-bot==13.15 requests telegraph tenacity