implemented to cancel to recursive download if some files failed to download.
This commit is contained in:
parent
c8b16939d7
commit
d392379235
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user