implemented mutiple download feature.

This commit is contained in:
Yoshihiro OKUMURA 2023-05-09 14:38:13 +09:00
parent c724af538b
commit e7197673fc
Signed by: orrisroot
GPG Key ID: 470AA444C92904B2
3 changed files with 72 additions and 17 deletions

View File

@ -80,3 +80,17 @@ class FileApi(BaseApi):
response = self._get(url) response = self._get(url)
self._raise_response_error(response) self._raise_response_error(response)
return response.json() 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

View File

@ -30,6 +30,12 @@ class UploadFile:
path: str path: str
@dataclass(frozen=True)
class DownloadFile:
file: File
path: str
class FileCommand(BaseCommand): class FileCommand(BaseCommand):
@staticmethod @staticmethod
def register(top_level_subparsers: _SubParsersAction) -> None: 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.add_argument("remote_path", help="Remote folder path (remote:/lab/path/)")
upload_parser.set_defaults(func=FileCommand.upload) upload_parser.set_defaults(func=FileCommand.upload)
# download # 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("remote_path", help="Remote file path (remote:/lab/path/file)")
download_parser.add_argument("local_path", help="Local folder path (/foo/bar/)") download_parser.add_argument("local_path", help="Local folder path (/foo/bar/)")
download_parser.set_defaults(func=FileCommand.download) download_parser.set_defaults(func=FileCommand.download)
@ -116,23 +125,28 @@ class FileCommand(BaseCommand):
parent_path = os.path.dirname(path) parent_path = os.path.dirname(path)
file_name = os.path.basename(path) file_name = os.path.basename(path)
connection = create_connection(remote) 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.") 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) laboratory = find_laboratory(connection, laboratory_name)
folder = find_folder(connection, laboratory, parent_path) parent_folder = find_folder(connection, laboratory, parent_path)
file = folder.find_file(file_name) file = parent_folder.find_file(file_name)
if file is None: download_files: list[DownloadFile] = []
raise IllegalArgumentException(f"File `{file_name}` not found.") if file is not None:
r = connection.session.get(connection.build_url("v2/" + file.download_url), stream=True) file_path = os.path.join(local_path, file_name)
try: download_files.append(DownloadFile(file, file_path))
with open(local_file, "wb") as f: else:
for chunk in r.iter_content(chunk_size=4096): if not args.recursive:
if chunk: raise IllegalArgumentException(f"Cannot download `{args.remote_path}`: Is a folder.")
f.write(chunk) sub_folder = parent_folder.find_sub_folder(file_name)
f.flush() if sub_folder is None:
except PermissionError: raise IllegalArgumentException(f"File or Folder`{file_name}` not found.")
raise IllegalArgumentException(f"Cannot create file `{local_file}`: Permission denied.") 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 @staticmethod
def move(args: Namespace) -> None: def move(args: Namespace) -> None:
@ -218,3 +232,30 @@ class FileCommand(BaseCommand):
pass pass
except MDRSException as e: except MDRSException as e:
print(f"API Error: {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)

View File

@ -2,7 +2,7 @@ from typing import Final
from pydantic.dataclasses import dataclass 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 from mdrsclient.models.utils import iso8601_to_user_friendly
ACCESS_LEVEL_NAMES: Final[dict[int, str]] = { ACCESS_LEVEL_NAMES: Final[dict[int, str]] = {