diff --git a/mdrsclient/api/file.py b/mdrsclient/api/file.py index ebbd34f..5655459 100644 --- a/mdrsclient/api/file.py +++ b/mdrsclient/api/file.py @@ -80,3 +80,17 @@ class FileApi(BaseApi): response = self._get(url) self._raise_response_error(response) return response.json() + + def download(self, file: File, path: str) -> bool: + print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name) + url = self.connection.build_url("v2/", file.download_url) + r = self.connection.session.get(url, stream=True) + try: + with open(path, "wb") as f: + for chunk in r.iter_content(chunk_size=4096): + if chunk: + f.write(chunk) + f.flush() + except PermissionError: + print(f"Cannot create file `{path}`: Permission denied.") + return True diff --git a/mdrsclient/commands/file.py b/mdrsclient/commands/file.py index 3520f20..3e73b01 100644 --- a/mdrsclient/commands/file.py +++ b/mdrsclient/commands/file.py @@ -30,6 +30,12 @@ class UploadFile: path: str +@dataclass(frozen=True) +class DownloadFile: + file: File + path: str + + class FileCommand(BaseCommand): @staticmethod def register(top_level_subparsers: _SubParsersAction) -> None: @@ -42,7 +48,10 @@ class FileCommand(BaseCommand): upload_parser.add_argument("remote_path", help="Remote folder path (remote:/lab/path/)") upload_parser.set_defaults(func=FileCommand.upload) # download - download_parser = top_level_subparsers.add_parser("download", help="download a file") + download_parser = top_level_subparsers.add_parser("download", help="download the file or folders") + download_parser.add_argument( + "-r", "--recursive", help="Download folders and their contents recursive", action="store_true" + ) download_parser.add_argument("remote_path", help="Remote file path (remote:/lab/path/file)") download_parser.add_argument("local_path", help="Local folder path (/foo/bar/)") download_parser.set_defaults(func=FileCommand.download) @@ -116,23 +125,28 @@ class FileCommand(BaseCommand): parent_path = os.path.dirname(path) file_name = os.path.basename(path) connection = create_connection(remote) - if not os.path.isdir(args.local_path): + local_path = os.path.abspath(args.local_path) + if not os.path.isdir(local_path): raise IllegalArgumentException(f"Local directory `{args.local_path}` not found.") - local_file = os.path.join(args.local_path, file_name) laboratory = find_laboratory(connection, laboratory_name) - folder = find_folder(connection, laboratory, parent_path) - file = folder.find_file(file_name) - if file is None: - raise IllegalArgumentException(f"File `{file_name}` not found.") - r = connection.session.get(connection.build_url("v2/" + file.download_url), stream=True) - try: - with open(local_file, "wb") as f: - for chunk in r.iter_content(chunk_size=4096): - if chunk: - f.write(chunk) - f.flush() - except PermissionError: - raise IllegalArgumentException(f"Cannot create file `{local_file}`: Permission denied.") + parent_folder = find_folder(connection, laboratory, parent_path) + file = parent_folder.find_file(file_name) + download_files: list[DownloadFile] = [] + if file is not None: + file_path = os.path.join(local_path, file_name) + download_files.append(DownloadFile(file, file_path)) + else: + if not args.recursive: + raise IllegalArgumentException(f"Cannot download `{args.remote_path}`: Is a folder.") + sub_folder = parent_folder.find_sub_folder(file_name) + if sub_folder is None: + raise IllegalArgumentException(f"File or Folder`{file_name}` not found.") + folder_api = FolderApi(connection) + sub_folder_dirname = os.path.join(local_path, sub_folder.name) + FileCommand._multiple_download_pickup_recursive_files( + folder_api, download_files, sub_folder.id, sub_folder_dirname + ) + FileCommand._multiple_download(connection, download_files) @staticmethod def move(args: Namespace) -> None: @@ -218,3 +232,30 @@ class FileCommand(BaseCommand): pass except MDRSException as e: print(f"API Error: {e}") + + @staticmethod + def _multiple_download_pickup_recursive_files( + folder_api: FolderApi, download_files: list[DownloadFile], folder_id: str, local_dirname: str + ) -> None: + folder = folder_api.retrieve(folder_id) + file_dirname = os.path.join(local_dirname, folder.name) + if not os.path.exists(file_dirname): + os.makedirs(file_dirname) + for file in folder.files: + file_path = os.path.join(file_dirname, file.name) + download_files.append(DownloadFile(file, file_path)) + for sub_folder in folder.sub_folders: + sub_folder_dirname = os.path.join(local_dirname, sub_folder.name) + FileCommand._multiple_download_pickup_recursive_files( + folder_api, download_files, sub_folder.id, sub_folder_dirname + ) + + @staticmethod + def _multiple_download(connection: MDRSConnection, download_files: list[DownloadFile]) -> None: + file_api = FileApi(connection) + with ThreadPoolExecutor(max_workers=NUMBER_OF_PROCESS) as pool: + pool.map(lambda x: FileCommand._multiple_download_worker(file_api, x), download_files) + + @staticmethod + def _multiple_download_worker(file_api: FileApi, download_file: DownloadFile) -> None: + file_api.download(download_file.file, download_file.path) diff --git a/mdrsclient/models/folder.py b/mdrsclient/models/folder.py index eebebb9..1c5fe3d 100644 --- a/mdrsclient/models/folder.py +++ b/mdrsclient/models/folder.py @@ -2,7 +2,7 @@ from typing import Final from pydantic.dataclasses import dataclass -from mdrsclient.models import File +from mdrsclient.models.file import File from mdrsclient.models.utils import iso8601_to_user_friendly ACCESS_LEVEL_NAMES: Final[dict[int, str]] = {