Add archival and extraction progress in compress module

- This adds the ability to show the progress of archival and extraction tasks
- Also adds the ability to cancel archival or extraction task
- Tidy up

bot/__init__.py

@@ -69,9 +69,9 @@ def get_config(name: str):
 
 try:
     users = get_config('AUTHORIZED_CHATS')
-    users = users.split(" ")
+    users = users.split()
     for user in users:
-        AUTHORIZED_CHATS.add(int(user))
+        AUTHORIZED_CHATS.add(int(user.strip()))
 except:
     pass
 

bot/__main__.py

@@ -67,6 +67,7 @@ def restart(update, context):
     restart_message = sendMessage("<b>Restart in progress...</b>", context.bot, update.message)
     if Interval:
         Interval[0].cancel()
+        Interval.clear()
     clean_all()
     with open(".restartmsg", "w") as f:
         f.truncate(0)

bot/helper/download_utils/__init__.py

@@ -0,0 +1 @@
+

bot/helper/drive_utils/gdriveTools.py

@@ -9,7 +9,7 @@
 from urllib.parse import parse_qs, urlparse
 from random import randrange
 from tenacity import retry, wait_exponential, stop_after_attempt, \
-    retry_if_exception_type, before_log, RetryError
+    retry_if_exception_type, RetryError
 
 from telegram import InlineKeyboardMarkup
 from telegraph.exceptions import RetryAfterError
@@ -18,7 +18,7 @@
 from google.oauth2 import service_account
 from google.oauth2.credentials import Credentials
 from googleapiclient.discovery import build
-from googleapiclient.errors import HttpError
+from googleapiclient.errors import Error as GCError, HttpError
 from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload
 
 from bot import LOGGER, DOWNLOAD_DIR, DRIVE_NAMES, DRIVE_IDS, INDEX_URLS, PARENT_ID, \
@@ -80,7 +80,7 @@ def speed(self):
         """
         try:
             return self.uploaded_bytes / self.total_time
-        except ZeroDivisionError:
+        except:
             return 0
 
     def dspeed(self):
@@ -90,7 +90,7 @@ def dspeed(self):
         """
         try:
             return self.downloaded_bytes / self.dtotal_time
-        except ZeroDivisionError:
+        except:
             return 0
 
     def cspeed(self):
@@ -100,7 +100,7 @@ def cspeed(self):
         """
         try:
             return self.transferred_size / int(time.time() - self.start_time)
-        except ZeroDivisionError:
+        except:
             return 0
 
     def authorize(self):
@@ -202,7 +202,7 @@ def __set_permission_email(self, file_id, email):
                    body=permissions,
                    sendNotificationEmail=False).execute()
 
-    def setPerm(self, link, access):
+    def setPermission(self, link, access):
         try:
             file_id = self.getIdFromUrl(link)
         except (KeyError, IndexError):
@@ -226,16 +226,15 @@ def setPerm(self, link, access):
                 token_service = self.alt_authorize()
                 if token_service is not None:
                     self.__service = token_service
-                    return self.setPerm(link, access)
+                    return self.setPermission(link, access)
                 msg = "Insufficient file permissions"
             else:
                 msg = str(err)
         return msg
 
     @retry(wait=wait_exponential(multiplier=2, min=3, max=6),
            stop=stop_after_attempt(3),
-           retry=retry_if_exception_type(HttpError),
-           before=before_log(LOGGER, logging.DEBUG))
+           retry=retry_if_exception_type(GCError))
     def copyFile(self, file_id, dest_id):
         body = {
             'parents': [dest_id]
@@ -266,8 +265,7 @@ def copyFile(self, file_id, dest_id):
 
     @retry(wait=wait_exponential(multiplier=2, min=3, max=6),
            stop=stop_after_attempt(3),
-           retry=retry_if_exception_type(HttpError),
-           before=before_log(LOGGER, logging.DEBUG))
+           retry=retry_if_exception_type(GCError))
     def getFileMetadata(self, file_id):
         return self.__service.files().get(
                    supportsAllDrives=True,
@@ -276,8 +274,7 @@ def getFileMetadata(self, file_id):
 
     @retry(wait=wait_exponential(multiplier=2, min=3, max=6),
            stop=stop_after_attempt(3),
-           retry=retry_if_exception_type(HttpError),
-           before=before_log(LOGGER, logging.DEBUG))
+           retry=retry_if_exception_type(GCError))
     def getFilesByFolderId(self, folder_id):
         page_token = None
         query = f"'{folder_id}' in parents and trashed = false"
@@ -357,10 +354,9 @@ def clone(self, link, key):
                 LOGGER.info(f"Total attempts: {err.last_attempt.attempt_number}")
                 err = err.last_attempt.exception()
             err = str(err).replace('>', '').replace('<', '')
-            LOGGER.error(err)
-            if "User rate limit exceeded" in str(err):
+            if "User rate limit exceeded" in err:
                 msg = "User rate limit exceeded"
-            elif "File not found" in str(err):
+            elif "File not found" in err:
                 token_service = self.alt_authorize()
                 if token_service is not None:
                     self.__service = token_service
@@ -389,8 +385,7 @@ def cloneFolder(self, name, local_path, folder_id, parent_id):
 
     @retry(wait=wait_exponential(multiplier=2, min=3, max=6),
            stop=stop_after_attempt(3),
-           retry=retry_if_exception_type(HttpError),
-           before=before_log(LOGGER, logging.DEBUG))
+           retry=retry_if_exception_type(GCError))
     def create_directory(self, directory_name, parent_id):
         file_metadata = {
             "name": directory_name,
@@ -437,8 +432,7 @@ def count(self, link):
                 LOGGER.info(f"Total attempts: {err.last_attempt.attempt_number}")
                 err = err.last_attempt.exception()
             err = str(err).replace('>', '').replace('<', '')
-            LOGGER.error(err)
-            if "File not found" in str(err):
+            if "File not found" in err:
                 token_service = self.alt_authorize()
                 if token_service is not None:
                     self.__service = token_service
@@ -493,8 +487,7 @@ def helper(self, link):
                 LOGGER.info(f"Total attempts: {err.last_attempt.attempt_number}")
                 err = err.last_attempt.exception()
             err = str(err).replace('>', '').replace('<', '')
-            LOGGER.error(err)
-            if "File not found" in str(err):
+            if "File not found" in err:
                 token_service = self.alt_authorize()
                 if token_service is not None:
                     self.__service = token_service
@@ -687,13 +680,10 @@ def upload(self, file_name: str):
                 link = f"https://drive.google.com/folderview?id={dir_id}"
                 if self.is_cancelled:
                     return
-        except Exception as e:
-            if isinstance(e, RetryError):
-                LOGGER.info(f"Total attempts: {e.last_attempt.attempt_number}")
-                err = e.last_attempt.exception()
-            else:
-                err = e
-            LOGGER.error(err)
+        except Exception as err:
+            if isinstance(err, RetryError):
+                LOGGER.info(f"Total attempts: {err.last_attempt.attempt_number}")
+                err = err.last_attempt.exception()
             self.__listener.onUploadError(str(err))
             self.is_errored = True
         finally:
@@ -732,8 +722,7 @@ def upload_dir(self, input_directory, parent_id):
 
     @retry(wait=wait_exponential(multiplier=2, min=3, max=6),
            stop=stop_after_attempt(3),
-           retry=(retry_if_exception_type(HttpError) | retry_if_exception_type(IOError)),
-           before=before_log(LOGGER, logging.DEBUG))
+           retry=(retry_if_exception_type(GCError) | retry_if_exception_type(IOError)))
     def upload_file(self, file_path, file_name, mime_type, parent_id):
         file_metadata = {
             'name': file_name,
@@ -788,10 +777,6 @@ def upload_file(self, file_path, file_name, mime_type, parent_id):
         download_url = self.__G_DRIVE_BASE_DOWNLOAD_URL.format(drive_file.get('id'))
         return download_url
 
-    @retry(wait=wait_exponential(multiplier=2, min=3, max=6),
-           stop=stop_after_attempt(3),
-           retry=retry_if_exception_type(HttpError),
-           before=before_log(LOGGER, logging.DEBUG))
     def _on_upload_progress(self):
         if self.status is not None:
             chunk_size = self.status.total_size * self.status.progress() - self._file_uploaded_bytes
@@ -816,10 +801,9 @@ def download(self, link):
                 LOGGER.info(f"Total attempts: {err.last_attempt.attempt_number}")
                 err = err.last_attempt.exception()
             err = str(err).replace('>', '').replace('<', '')
-            LOGGER.error(err)
-            if "downloadQuotaExceeded" in str(err):
+            if "downloadQuotaExceeded" in err:
                 err = "Download quota exceeded."
-            elif "File not found" in str(err):
+            elif "File not found" in err:
                 token_service = self.alt_authorize()
                 if token_service is not None:
                     self.__service = token_service
@@ -860,11 +844,15 @@ def download_folder(self, folder_id, path, folder_name):
 
     @retry(wait=wait_exponential(multiplier=2, min=3, max=6),
            stop=stop_after_attempt(3),
-           retry=(retry_if_exception_type(HttpError) | retry_if_exception_type(IOError)),
-           before=before_log(LOGGER, logging.DEBUG))
+           retry=(retry_if_exception_type(GCError) | retry_if_exception_type(IOError)))
     def download_file(self, file_id, path, filename, mime_type):
         request = self.__service.files().get_media(fileId=file_id)
         filename = filename.replace('/', '')
+        if len(filename.encode()) > 255:
+            ext = os.path.splitext(filename)[1]
+            filename = filename[:245] + ext
+            if self.name.endswith(ext):
+                self.name = filename
         fh = FileIO('{}{}'.format(path, filename), 'wb')
         downloader = MediaIoBaseDownload(fh, request, chunksize=50 * 1024 * 1024)
         done = False
@@ -891,10 +879,6 @@ def download_file(self, file_id, path, filename, mime_type):
                         raise err
         self._file_downloaded_bytes = 0
 
-    @retry(wait=wait_exponential(multiplier=2, min=3, max=6),
-           stop=stop_after_attempt(3),
-           retry=retry_if_exception_type(HttpError),
-           before=before_log(LOGGER, logging.DEBUG))
     def _on_download_progress(self):
         if self.dstatus is not None:
             chunk_size = self.dstatus.total_size * self.dstatus.progress() - self._file_downloaded_bytes

bot/helper/ext_utils/bot_utils.py

@@ -13,9 +13,9 @@
 URL_REGEX = r'(?:(?:https?|ftp):\/\/)?[\w/\-?=%.]+\.[\w/\-?=%.]+'
 
 class TaskStatus:
-    STATUS_UPLOADING = "Uploading...📤"
-    STATUS_DOWNLOADING = "Downloading...📥"
     STATUS_CLONING = "Cloning...♻️"
+    STATUS_DOWNLOADING = "Downloading...📥"
+    STATUS_UPLOADING = "Uploading...📤"
     STATUS_ARCHIVING = "Archiving...🔐"
     STATUS_EXTRACTING = "Extracting...📂"
 
@@ -39,11 +39,8 @@ def cancel(self):
 def getDownloadByGid(gid):
     with download_dict_lock:
         for dl in list(download_dict.values()):
-            status = dl.status()
-            if status not in [TaskStatus.STATUS_ARCHIVING,
-                              TaskStatus.STATUS_EXTRACTING]:
-                if dl.gid() == gid:
-                    return dl
+            if dl.gid() == gid:
+                return dl
     return None
 
 def get_progress_bar_string(status):
@@ -63,21 +60,23 @@ def get_readable_message():
         for download in list(download_dict.values()):
             msg += f"<b>Name:</b> <code>{escape(str(download.name()))}</code>"
             msg += f"\n<b>Status:</b> <i>{download.status()}</i>"
-            if download.status() not in [TaskStatus.STATUS_ARCHIVING,
-                                         TaskStatus.STATUS_EXTRACTING]:
-                msg += f"\n{get_progress_bar_string(download)} {download.progress()}"
-                if download.status() == TaskStatus.STATUS_CLONING:
-                    msg += f"\n<b>Cloned:</b> {get_readable_file_size(download.processed_bytes())} / {download.size()}"
-                    msg += f"\n<b>Transfers:</b> {download.processed_files()} / {download.files()}"
-                elif download.status() == TaskStatus.STATUS_UPLOADING:
-                    msg += f"\n<b>Uploaded:</b> {get_readable_file_size(download.processed_bytes())} / {download.size()}"
-                else:
-                    msg += f"\n<b>Downloaded:</b> {get_readable_file_size(download.processed_bytes())} / {download.size()}"
-                msg += f"\n<b>Speed:</b> {download.speed()} | <b>ETA:</b> {download.eta()}"
-                msg += f"\n<code>/{BotCommands.CancelCommand} {download.gid()}</code>"
-            else:
-                msg += f"\n<b>Size: </b>{download.size()}"
+            msg += f"\n{get_progress_bar_string(download)} {download.progress()}"
+            if download.status() == TaskStatus.STATUS_CLONING:
+                msg += f"\n<b>Cloned:</b> {get_readable_file_size(download.processed_bytes())} / {download.size()}"
+                msg += f"\n<b>Transfers:</b> {download.processed_files()} / {download.files()}"
+            elif download.status() == TaskStatus.STATUS_DOWNLOADING:
+                msg += f"\n<b>Downloaded:</b> {get_readable_file_size(download.processed_bytes())} / {download.size()}"
+            elif download.status() == TaskStatus.STATUS_UPLOADING:
+                msg += f"\n<b>Uploaded:</b> {get_readable_file_size(download.processed_bytes())} / {download.size()}"
+            elif download.status() == TaskStatus.STATUS_ARCHIVING:
+                msg += f"\n<b>Archived:</b> {get_readable_file_size(download.processed_bytes())} / {download.size()}"
+            elif download.status() == TaskStatus.STATUS_EXTRACTING:
+                msg += f"\n<b>Extracted:</b> {get_readable_file_size(download.processed_bytes())} / {download.size()}"
+            msg += f"\n<b>Speed:</b> {download.speed()} | <b>ETA:</b> {download.eta()}"
+            msg += f"\n<code>/{BotCommands.CancelCommand} {download.gid()}</code>"
             msg += "\n\n"
+        if len(msg) == 0:
+            return None
         cpu = cpu_percent(interval=0.5)
         ram = virtual_memory().percent
         disk = disk_usage('/').percent

bot/helper/ext_utils/database.py

@@ -24,5 +24,4 @@ def load_users(self):
             AUTHORIZED_CHATS.add(user['user_id'])
 
 if DATABASE_URL is not None:
-    db = DatabaseHelper()
-    db.load_users()
+    DatabaseHelper().load_users()