From ab6c01146e116e25a260f41b55a8e2efaf4c3081 Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Thu, 26 Sep 2019 18:21:07 +0530 Subject: [PATCH] Make separate module for fs related jobs Signed-off-by: lzzy12 --- bot/helper/fs_utils.py | 29 +++++++++++++++++++++++++++++ bot/helper/gdriveTools.py | 17 +++++++++-------- 2 files changed, 38 insertions(+), 8 deletions(-) create mode 100644 bot/helper/fs_utils.py diff --git a/bot/helper/fs_utils.py b/bot/helper/fs_utils.py new file mode 100644 index 0000000..3a950ee --- /dev/null +++ b/bot/helper/fs_utils.py @@ -0,0 +1,29 @@ +import sys +from bot import aria2, LOGGER, DOWNLOAD_DIR +import shutil +import os +import pathlib + + +def clean_download(path: str): + if os.path.exists(path): + shutil.rmtree(path) + + +def exit_clean_up(signal, frame): + try: + LOGGER.info("Please wait, while we clean up the downloads and stop running downloads") + aria2.autopurge() + aria2.remove_all(True) + shutil.rmtree(DOWNLOAD_DIR) + sys.exit(0) + except KeyboardInterrupt: + LOGGER.warning("Force Exiting before the cleanup finishes!") + sys.exit(1) + + +def tar(orig_path: str): + path = pathlib.PurePath(orig_path) + base = path.name + root = path.parent.name + shutil.make_archive(base_name=orig_path, format='tar', root_dir=root, base_dir=base) diff --git a/bot/helper/gdriveTools.py b/bot/helper/gdriveTools.py index 5602760..1f6a0e6 100644 --- a/bot/helper/gdriveTools.py +++ b/bot/helper/gdriveTools.py @@ -7,7 +7,7 @@ from mimetypes import guess_type import os from bot import LOGGER, parent_id, DOWNLOAD_DIR, download_list from .listeners import MirrorListeners -from shutil import rmtree +from .fs_utils import clean_download class GoogleDriveHelper: @@ -76,22 +76,23 @@ class GoogleDriveHelper: else: file_name, mime_type = self.file_ops(file_path) try: - dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_id) - self.upload_dir(file_path, dir_id) + dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name))) + self.upload_dir(file_path) LOGGER.info("Uploaded To G-Drive: " + file_name) link = "https://drive.google.com/folderview?id={}".format(dir_id) except Exception as e: LOGGER.error(str(e)) self.__listener.onUploadError(str(e)) + clean_download(file_dir) raise Exception('Error: {}'.format(str(e))) del download_list[self.__listener.update.update_id] LOGGER.info(download_list) self.__listener.onUploadComplete(link, file_name) LOGGER.info("Deleting downloaded file/folder..") - rmtree(file_dir) + clean_download(file_dir) return link - def create_directory(self, directory_name, parent_id): + def create_directory(self, directory_name): permissions = { "role": "reader", "type": "anyone", @@ -110,7 +111,7 @@ class GoogleDriveHelper: LOGGER.info("Created Google-Drive Folder:\nName: {}\nID: {} ".format(file.get("title"), file_id)) return file_id - def upload_dir(self, input_directory, parent_id): + def upload_dir(self, input_directory): list_dirs = os.listdir(input_directory) if len(list_dirs) == 0: return parent_id @@ -118,8 +119,8 @@ class GoogleDriveHelper: for a_c_f_name in list_dirs: current_file_name = os.path.join(input_directory, a_c_f_name) if os.path.isdir(current_file_name): - current_dir_id = self.create_directory(a_c_f_name, parent_id) - r_p_id = self.upload_dir(current_file_name, current_dir_id) + current_dir_id = self.create_directory(a_c_f_name) + r_p_id = self.upload_dir(current_file_name) else: file_name, mime_type = self.file_ops(current_file_name) # current_file_name will have the full path