diff --git a/mdrsclient/commands/download.py b/mdrsclient/commands/download.py index 028f122..d887630 100644 --- a/mdrsclient/commands/download.py +++ b/mdrsclient/commands/download.py @@ -8,7 +8,7 @@ from pydantic.dataclasses import dataclass from mdrsclient.api import FilesApi, FoldersApi from mdrsclient.commands.base import BaseCommand from mdrsclient.connection import MDRSConnection -from mdrsclient.exceptions import IllegalArgumentException +from mdrsclient.exceptions import IllegalArgumentException, UnexpectedException from mdrsclient.models import File from mdrsclient.settings import CONCURRENT @@ -19,6 +19,12 @@ class DownloadFileInfo: path: str +@dataclass +class DownloadContext: + hasError: bool + files: list[DownloadFileInfo] + + class DownloadCommand(BaseCommand): @classmethod def register(cls, parsers: Any) -> None: @@ -52,10 +58,11 @@ class DownloadCommand(BaseCommand): laboratory = cls._find_laboratory(connection, laboratory_name) r_parent_folder = cls._find_folder(connection, laboratory, r_dirname, password) file = r_parent_folder.find_file(r_basename) - download_files: list[DownloadFileInfo] = [] if file is not None: + context = DownloadContext(False, []) l_path = os.path.join(l_dirname, r_basename) - download_files.append(DownloadFileInfo(file, l_path)) + context.files.append(DownloadFileInfo(file, l_path)) + cls.__multiple_download(connection, context) else: folder = r_parent_folder.find_sub_folder(r_basename) if folder is None: @@ -63,13 +70,13 @@ class DownloadCommand(BaseCommand): if not is_recursive: raise IllegalArgumentException(f"Cannot download `{r_path}`: Is a folder.") folder_api = FoldersApi(connection) - cls.__multiple_download_pickup_recursive_files(folder_api, download_files, folder.id, l_dirname) - cls.__multiple_download(connection, download_files) + cls.__multiple_download_pickup_recursive_files(connection, folder_api, folder.id, l_dirname) @classmethod def __multiple_download_pickup_recursive_files( - cls, folder_api: FoldersApi, infolist: list[DownloadFileInfo], folder_id: str, basedir: str + cls, connection: MDRSConnection, folder_api: FoldersApi, folder_id: str, basedir: str ) -> None: + context = DownloadContext(False, []) folder = folder_api.retrieve(folder_id) dirname = os.path.join(basedir, folder.name) if not os.path.exists(dirname): @@ -77,17 +84,27 @@ class DownloadCommand(BaseCommand): print(dirname) for file in folder.files: path = os.path.join(dirname, file.name) - infolist.append(DownloadFileInfo(file, path)) + context.files.append(DownloadFileInfo(file, path)) + cls.__multiple_download(connection, context) + if context.hasError: + raise UnexpectedException("Some files failed to download.") for sub_folder in folder.sub_folders: - cls.__multiple_download_pickup_recursive_files(folder_api, infolist, sub_folder.id, dirname) + cls.__multiple_download_pickup_recursive_files(connection, folder_api, sub_folder.id, dirname) @classmethod - def __multiple_download(cls, connection: MDRSConnection, infolist: list[DownloadFileInfo]) -> None: + def __multiple_download(cls, connection: MDRSConnection, context: DownloadContext) -> None: file_api = FilesApi(connection) with ThreadPoolExecutor(max_workers=CONCURRENT) as pool: - pool.map(lambda x: cls.__multiple_download_worker(file_api, x), infolist) + results = pool.map(lambda x: cls.__multiple_download_worker(file_api, x), 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) -> None: - file_api.download(info.file, info.path) + def __multiple_download_worker(cls, file_api: FilesApi, info: DownloadFileInfo) -> bool: + try: + file_api.download(info.file, info.path) + except Exception: + return False print(info.path) + return True