implemented to cancel to recursive download if some files failed to download.

This commit is contained in:
Yoshihiro OKUMURA 2024-07-08 20:35:48 +09:00
parent c8b16939d7
commit d392379235
Signed by: orrisroot
GPG Key ID: 470AA444C92904B2

View File

@ -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:
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