94 lines
4.1 KiB
Python
94 lines
4.1 KiB
Python
import os
|
|
from argparse import Namespace
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import Any
|
|
|
|
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.models import File
|
|
from mdrsclient.settings import CONCURRENT
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DownloadFileInfo:
|
|
file: File
|
|
path: str
|
|
|
|
|
|
class DownloadCommand(BaseCommand):
|
|
@classmethod
|
|
def register(cls, parsers: Any) -> None:
|
|
download_parser = parsers.add_parser("download", help="download the file or folder")
|
|
download_parser.add_argument(
|
|
"-r", "--recursive", help="download folders and their contents recursive", action="store_true"
|
|
)
|
|
download_parser.add_argument("-p", "--password", help="password to use when open locked folder")
|
|
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=cls.func)
|
|
|
|
@classmethod
|
|
def func(cls, args: Namespace) -> None:
|
|
remote_path = str(args.remote_path)
|
|
local_path = str(args.local_path)
|
|
is_recursive = bool(args.recursive)
|
|
password = str(args.password) if args.password else None
|
|
cls.download(remote_path, local_path, is_recursive, password)
|
|
|
|
@classmethod
|
|
def download(cls, remote_path: str, local_path: str, is_recursive: bool, password: str | None) -> None:
|
|
(remote, laboratory_name, r_path) = cls._parse_remote_host_with_path(remote_path)
|
|
r_path = r_path.rstrip("/")
|
|
r_dirname = os.path.dirname(r_path)
|
|
r_basename = os.path.basename(r_path)
|
|
connection = cls._create_connection(remote)
|
|
l_dirname = os.path.realpath(local_path)
|
|
if not os.path.isdir(l_dirname):
|
|
raise IllegalArgumentException(f"Local directory `{local_path}` not found.")
|
|
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:
|
|
l_path = os.path.join(l_dirname, r_basename)
|
|
download_files.append(DownloadFileInfo(file, l_path))
|
|
else:
|
|
folder = r_parent_folder.find_sub_folder(r_basename)
|
|
if folder is None:
|
|
raise IllegalArgumentException(f"File or folder `{r_path}` not found.")
|
|
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)
|
|
|
|
@classmethod
|
|
def __multiple_download_pickup_recursive_files(
|
|
cls, folder_api: FoldersApi, infolist: list[DownloadFileInfo], folder_id: str, basedir: str
|
|
) -> None:
|
|
folder = folder_api.retrieve(folder_id)
|
|
dirname = os.path.join(basedir, folder.name)
|
|
if not os.path.exists(dirname):
|
|
os.makedirs(dirname)
|
|
print(dirname)
|
|
for file in folder.files:
|
|
path = os.path.join(dirname, file.name)
|
|
infolist.append(DownloadFileInfo(file, path))
|
|
for sub_folder in folder.sub_folders:
|
|
cls.__multiple_download_pickup_recursive_files(folder_api, infolist, sub_folder.id, dirname)
|
|
|
|
@classmethod
|
|
def __multiple_download(cls, connection: MDRSConnection, infolist: list[DownloadFileInfo]) -> None:
|
|
file_api = FilesApi(connection)
|
|
with ThreadPoolExecutor(max_workers=CONCURRENT) as pool:
|
|
pool.map(lambda x: cls.__multiple_download_worker(file_api, x), infolist)
|
|
|
|
@classmethod
|
|
def __multiple_download_worker(cls, file_api: FilesApi, info: DownloadFileInfo) -> None:
|
|
file_api.download(info.file, info.path)
|
|
print(info.path)
|