From 24172dc65c5ca680ffa47846da904ad969eddd63 Mon Sep 17 00:00:00 2001
From: Yoshihiro OKUMURA <orrisroot@gmail.com>
Date: Wed, 18 Sep 2024 10:56:40 +0900
Subject: [PATCH] implemented -s --skip-if-file-exists option for download
 commnad.

---
 mdrsclient/commands/download.py | 37 ++++++++++++++++++++++-----------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/mdrsclient/commands/download.py b/mdrsclient/commands/download.py
index 3aab291..c17bf71 100644
--- a/mdrsclient/commands/download.py
+++ b/mdrsclient/commands/download.py
@@ -22,6 +22,7 @@ class DownloadFileInfo:
 @dataclass
 class DownloadContext:
     hasError: bool
+    isSkipIfExists: bool
     files: list[DownloadFileInfo]
 
 
@@ -32,6 +33,12 @@ class DownloadCommand(BaseCommand):
         download_parser.add_argument(
             "-r", "--recursive", help="download folders and their contents recursive", action="store_true"
         )
+        download_parser.add_argument(
+            "-s",
+            "--skip-if-exists",
+            help="skip the download if file is already downloaded and file size is the same",
+            action="store_true",
+        )
         download_parser.add_argument(
             "-e", "--exclude", help="exclude to download path matched file or folders", action="append"
         )
@@ -45,13 +52,15 @@ class DownloadCommand(BaseCommand):
         remote_path = str(args.remote_path)
         local_path = str(args.local_path)
         is_recursive = bool(args.recursive)
+        is_skip_if_exists = bool(args.skip_if_exists)
         password = str(args.password) if args.password else None
         excludes = list(map(lambda x: str(x).rstrip("/").lower(), args.exclude)) if args.exclude is not None else []
-        cls.download(remote_path, local_path, is_recursive, password, excludes)
+        cls.download(remote_path, local_path, is_recursive, is_skip_if_exists, password, excludes)
 
     @classmethod
     def download(
-        cls, remote_path: str, local_path: str, is_recursive: bool, password: str | None, excludes: list[str]
+        cls, remote_path: str, local_path: str, is_recursive: bool, is_skip_if_exists: bool, password: str | None,
+        excludes: list[str]
     ) -> None:
         (remote, laboratory_name, r_path) = cls._parse_remote_host_with_path(remote_path)
         r_path = r_path.rstrip("/")
@@ -67,7 +76,7 @@ class DownloadCommand(BaseCommand):
         if file is not None:
             if cls.__check_excludes(excludes, laboratory, r_parent_folder, file):
                 return
-            context = DownloadContext(False, [])
+            context = DownloadContext(False, is_skip_if_exists, [])
             l_path = os.path.join(l_dirname, r_basename)
             context.files.append(DownloadFileInfo(file, l_path))
             cls.__multiple_download(connection, context)
@@ -79,7 +88,7 @@ class DownloadCommand(BaseCommand):
                 raise IllegalArgumentException(f"Cannot download `{r_path}`: Is a folder.")
             folder_api = FoldersApi(connection)
             cls.__multiple_download_pickup_recursive_files(
-                connection, folder_api, laboratory, folder.id, l_dirname, excludes
+                connection, folder_api, laboratory, folder.id, l_dirname, excludes, is_skip_if_exists
             )
 
     @classmethod
@@ -91,8 +100,9 @@ class DownloadCommand(BaseCommand):
         folder_id: str,
         basedir: str,
         excludes: list[str],
+        is_skip_if_exists: bool
     ) -> None:
-        context = DownloadContext(False, [])
+        context = DownloadContext(False, is_skip_if_exists, [])
         folder = folder_api.retrieve(folder_id)
         dirname = os.path.join(basedir, folder.name)
         if cls.__check_excludes(excludes, laboratory, folder, None):
@@ -110,24 +120,27 @@ class DownloadCommand(BaseCommand):
             raise UnexpectedException("Some files failed to download.")
         for sub_folder in folder.sub_folders:
             cls.__multiple_download_pickup_recursive_files(
-                connection, folder_api, laboratory, sub_folder.id, dirname, excludes
+                connection, folder_api, laboratory, sub_folder.id, dirname, excludes, is_skip_if_exists
             )
 
     @classmethod
     def __multiple_download(cls, connection: MDRSConnection, context: DownloadContext) -> None:
         file_api = FilesApi(connection)
         with ThreadPoolExecutor(max_workers=CONCURRENT) as pool:
-            results = pool.map(lambda x: cls.__multiple_download_worker(file_api, x), context.files)
+            results = pool.map(
+                lambda x: cls.__multiple_download_worker(file_api, x, context.isSkipIfExists), context.files
+            )
             hasError = next(filter(lambda x: x is False, results), None)
             if hasError is not None:
                 context.hasError = True
 
     @classmethod
-    def __multiple_download_worker(cls, file_api: FilesApi, info: DownloadFileInfo) -> bool:
-        try:
-            file_api.download(info.file, info.path)
-        except Exception:
-            return False
+    def __multiple_download_worker(cls, file_api: FilesApi, info: DownloadFileInfo, is_skip_if_exists: bool) -> bool:
+        if not is_skip_if_exists or os.path.exists(info.path) and info.file.size != os.path.getsize(info.path):
+            try:
+                file_api.download(info.file, info.path)
+            except Exception:
+                return False
         print(info.path)
         return True