diff --git a/datashuttle/configs/config_class.py b/datashuttle/configs/config_class.py index 5bf203afc..8bac7faf9 100644 --- a/datashuttle/configs/config_class.py +++ b/datashuttle/configs/config_class.py @@ -250,7 +250,7 @@ def get_base_folder( def get_rclone_config_name( self, connection_method: Optional[str] = None ) -> str: - """Generate the rclone configuration name for the project. + """Generate the rclone configuration name for the central project. These configs are created by datashuttle but managed and stored by rclone. """ diff --git a/datashuttle/datashuttle_class.py b/datashuttle/datashuttle_class.py index f78c64347..a0496fcce 100644 --- a/datashuttle/datashuttle_class.py +++ b/datashuttle/datashuttle_class.py @@ -845,7 +845,8 @@ def setup_google_drive_connection(self) -> None: secret if `gdrive_client_id` is set in the configs. Next, the user will be asked if their machine has access to a browser. - If not, they will be prompted to input their service account file path. + If not, they will be prompted to input a config_token after running an + rclone command displayed to the user on a machine with access to a browser. Next, with the provided credentials, the final setup will be done. This opens up a browser if the user confirmed access to a browser. @@ -855,23 +856,25 @@ def setup_google_drive_connection(self) -> None: local_vars=locals(), ) - browser_available = gdrive.ask_user_for_browser(log=True) - - service_account_filepath = None - gdrive_client_secret = None - - if browser_available and self.cfg["gdrive_client_id"]: + if self.cfg["gdrive_client_id"]: gdrive_client_secret = gdrive.get_client_secret() + else: + gdrive_client_secret = None - elif not browser_available: - service_account_filepath = ( - gdrive.prompt_and_get_service_account_filepath( - log=True, - ) + browser_available = gdrive.ask_user_for_browser(log=True) + + if not browser_available: + config_token = gdrive.prompt_and_get_config_token( + self.cfg, + gdrive_client_secret, + self.cfg.get_rclone_config_name("gdrive"), + log=True, ) + else: + config_token = None process = self._setup_rclone_gdrive_config( - gdrive_client_secret, service_account_filepath + gdrive_client_secret, config_token ) rclone.await_call_rclone_with_popen_raise_on_fail(process, log=True) @@ -1527,13 +1530,13 @@ def _setup_rclone_central_local_filesystem_config(self) -> None: def _setup_rclone_gdrive_config( self, gdrive_client_secret: str | None, - service_account_filepath: str | None, + config_token: str | None, ) -> subprocess.Popen: return rclone.setup_rclone_config_for_gdrive( self.cfg, self.cfg.get_rclone_config_name("gdrive"), gdrive_client_secret, - service_account_filepath, + config_token, ) def _setup_rclone_aws_config( diff --git a/datashuttle/tui/interface.py b/datashuttle/tui/interface.py index 05dd84035..82c2d0d97 100644 --- a/datashuttle/tui/interface.py +++ b/datashuttle/tui/interface.py @@ -13,7 +13,7 @@ from datashuttle import DataShuttle from datashuttle.configs import load_configs -from datashuttle.utils import aws, rclone, ssh, utils +from datashuttle.utils import aws, gdrive, rclone, ssh, utils class Interface: @@ -511,7 +511,7 @@ def setup_key_pair_and_rclone_config( def setup_google_drive_connection( self, gdrive_client_secret: Optional[str] = None, - service_account_filepath: Optional[str] = None, + config_token: Optional[str] = None, ) -> InterfaceOutput: """Try to set up and validate connection to Google Drive. @@ -526,7 +526,7 @@ def setup_google_drive_connection( """ try: process = self.project._setup_rclone_gdrive_config( - gdrive_client_secret, service_account_filepath + gdrive_client_secret, config_token ) self.google_drive_rclone_setup_process = process self.gdrive_setup_process_killed = False @@ -539,6 +539,21 @@ def setup_google_drive_connection( except BaseException as e: return False, str(e) + def get_rclone_message_for_gdrive_without_browser( + self, gdrive_client_secret: Optional[str] = None + ) -> InterfaceOutput: + """Get the rclone message for Google Drive setup without a browser.""" + try: + output = gdrive.preliminary_for_setup_without_browser( + self.project.cfg, + gdrive_client_secret, + self.project.cfg.get_rclone_config_name("gdrive"), + log=False, + ) + return True, output + except BaseException as e: + return False, str(e) + def terminate_google_drive_setup(self) -> None: """Terminate rclone setup for google drive by killing the rclone process.""" assert self.google_drive_rclone_setup_process is not None diff --git a/datashuttle/tui/screens/setup_gdrive.py b/datashuttle/tui/screens/setup_gdrive.py index 3f04e0c3b..0ffb5eba4 100644 --- a/datashuttle/tui/screens/setup_gdrive.py +++ b/datashuttle/tui/screens/setup_gdrive.py @@ -25,8 +25,8 @@ class SetupGdriveScreen(ModalScreen): If the config contains a "gdrive_client_id", the user is prompted to enter a client secret. If the user has access to a browser, a Google Drive - authentication page will open. Otherwise, the user is asked to enter path to a - service account file. + authentication page will open. Otherwise, the user is asked to run an rclone command + and input a config token. """ def __init__(self, interface: Interface) -> None: @@ -37,6 +37,7 @@ def __init__(self, interface: Interface) -> None: self.stage: int = 0 self.setup_worker: Worker | None = None self.is_browser_available: bool = True + self.gdrive_client_secret: Optional[str] = None # For handling credential inputs self.input_box: Input = Input( @@ -75,9 +76,9 @@ def on_button_pressed(self, event: Button.Pressed) -> None: and proceeds to a browser authentication. 3) "no" button : A "no" answer to the availability of browser question. On click, - prompts the user to enter path to their service account file. + prompts the user to enter a config token by running an rclone command. - 4) "enter" button : To enter the client secret or service account file path. + 4) "enter" button : To enter the client secret or config token. 5) "finish" button : To finish the setup. @@ -94,39 +95,44 @@ def on_button_pressed(self, event: Button.Pressed) -> None: self.dismiss() elif event.button.id == "setup_gdrive_ok_button": - self.ask_user_for_browser() - - elif event.button.id == "setup_gdrive_yes_button": - self.remove_yes_no_buttons() + self.query_one("#setup_gdrive_ok_button").remove() if self.interface.project.cfg["gdrive_client_id"]: self.ask_user_for_gdrive_client_secret() else: - self.open_browser_and_setup_gdrive_connection() + self.ask_user_for_browser() + + elif event.button.id == "setup_gdrive_yes_button": + self.remove_yes_no_buttons() + self.open_browser_and_setup_gdrive_connection( + self.gdrive_client_secret + ) elif event.button.id == "setup_gdrive_no_button": self.is_browser_available = False self.remove_yes_no_buttons() - self.prompt_user_for_service_account_filepath() + self.prompt_user_for_config_token() elif event.button.id == "setup_gdrive_enter_button": - if self.is_browser_available: - gdrive_client_secret = ( + if ( + self.interface.project.cfg["gdrive_client_id"] + and self.stage == 0 + ): + self.gdrive_client_secret = ( self.input_box.value.strip() if self.input_box.value.strip() else None ) - self.open_browser_and_setup_gdrive_connection( - gdrive_client_secret - ) + self.stage += 1 + self.ask_user_for_browser() else: - service_account_filepath = ( + config_token = ( self.input_box.value.strip() if self.input_box.value.strip() else None ) - self.setup_gdrive_connection_using_service_account( - service_account_filepath + self.setup_gdrive_connection_using_config_token( + self.gdrive_client_secret, config_token ) def ask_user_for_browser(self) -> None: @@ -137,7 +143,11 @@ def ask_user_for_browser(self) -> None: ) self.update_message_box_message(message) - self.query_one("#setup_gdrive_ok_button").remove() + if self.enter_button.is_mounted: + self.enter_button.remove() + + if self.input_box.is_mounted: + self.input_box.visible = False # Mount the Yes and No buttons yes_button = Button("Yes", id="setup_gdrive_yes_button") @@ -184,34 +194,46 @@ def open_browser_and_setup_gdrive_connection( ) ) - def prompt_user_for_service_account_filepath(self) -> None: - """Set up widgets and prompt user for their service account file path for browser-less connection.""" - message = "Please enter your service account file path." + def prompt_user_for_config_token(self) -> None: + """Prompt the user for the rclone config token for Google Drive setup.""" + success, message = ( + self.interface.get_rclone_message_for_gdrive_without_browser( + self.gdrive_client_secret + ) + ) - self.update_message_box_message(message) + if not success: + self.display_failed(message) + return + self.update_message_box_message( + message + "\nPress shift+click to copy." + ) + + self.enter_button = Button("Enter", id="setup_gdrive_enter_button") self.query_one("#setup_gdrive_buttons_horizontal").mount( self.enter_button, before="#setup_gdrive_cancel_button" ) self.mount_input_box_before_buttons() - def setup_gdrive_connection_using_service_account( - self, service_account_filepath: Optional[str] = None + def setup_gdrive_connection_using_config_token( + self, gdrive_client_secret: str | None, config_token: str | None ) -> None: - """Set up the Google Drive connection using service account and show success message.""" + """Set up the Google Drive connection using rclone config token.""" message = "Setting up connection." self.update_message_box_message(message) asyncio.create_task( self.setup_gdrive_connection_and_update_ui( - service_account_filepath=service_account_filepath + gdrive_client_secret=gdrive_client_secret, + config_token=config_token, ) ) async def setup_gdrive_connection_and_update_ui( self, gdrive_client_secret: Optional[str] = None, - service_account_filepath: Optional[str] = None, + config_token: Optional[str] = None, ) -> None: """Start the Google Drive connection setup in a separate thread. @@ -232,7 +254,7 @@ async def setup_gdrive_connection_and_update_ui( self.enter_button.disabled = True worker = self.setup_gdrive_connection( - gdrive_client_secret, service_account_filepath + gdrive_client_secret, config_token ) self.setup_worker = worker if worker.is_running: @@ -250,7 +272,7 @@ async def setup_gdrive_connection_and_update_ui( def setup_gdrive_connection( self, gdrive_client_secret: Optional[str] = None, - service_account_filepath: Optional[str] = None, + config_token: Optional[str] = None, ) -> Worker[InterfaceOutput]: """Authenticate the Google Drive connection. @@ -260,7 +282,7 @@ def setup_gdrive_connection( with Google Drive. """ success, output = self.interface.setup_google_drive_connection( - gdrive_client_secret, service_account_filepath + gdrive_client_secret, config_token ) return success, output @@ -302,6 +324,8 @@ def mount_input_box_before_buttons( self.query_one("#gdrive_setup_messagebox_message_container").mount( self.input_box, after="#gdrive_setup_messagebox_message" ) + self.input_box.visible = True + self.input_box.value = "" def remove_yes_no_buttons(self) -> None: """Remove yes and no buttons.""" diff --git a/datashuttle/tui/shared/configs_content.py b/datashuttle/tui/shared/configs_content.py index 88ab752bf..70410d7d2 100644 --- a/datashuttle/tui/shared/configs_content.py +++ b/datashuttle/tui/shared/configs_content.py @@ -530,7 +530,7 @@ def widget_configs_match_saved_configs(self): for key, value in cfg_kwargs.items(): saved_val = self.interface.get_configs()[key] if key in ["central_path", "local_path"]: - if value.name != project_name: + if value is not None and value.name != project_name: value = value / project_name if saved_val != value: return False diff --git a/datashuttle/utils/folders.py b/datashuttle/utils/folders.py index c5abf696f..74194660b 100644 --- a/datashuttle/utils/folders.py +++ b/datashuttle/utils/folders.py @@ -17,11 +17,11 @@ from datashuttle.configs.config_class import Configs from datashuttle.utils.custom_types import TopLevelFolder -import glob +import fnmatch from pathlib import Path from datashuttle.configs import canonical_folders, canonical_tags -from datashuttle.utils import rclone, ssh, utils, validation +from datashuttle.utils import rclone, utils, validation from datashuttle.utils.custom_exceptions import NeuroBlueprintError # ----------------------------------------------------------------------------- @@ -599,56 +599,62 @@ def search_for_folders( Discovered folders (`all_folder_names`) and files (`all_filenames`). """ - if local_or_central == "central" and cfg["connection_method"] in [ - "ssh", - "gdrive", - "aws", - ]: - if cfg["connection_method"] == "ssh": - all_folder_names, all_filenames = ( - ssh.search_ssh_central_for_folders( - search_path, - search_prefix, - cfg, - verbose, - return_full_path, - ) - ) - - else: - all_folder_names, all_filenames = search_gdrive_or_aws_for_folders( - search_path, search_prefix, cfg, return_full_path - ) - + if ( + local_or_central == "local" + or cfg["connection_method"] == "local_filesystem" + ) and not search_path.exists(): + if verbose: + utils.log_and_message(f"No file found at {search_path.as_posix()}") + return [], [] + + if local_or_central == "local": + rclone_config_name = None else: - if not search_path.exists(): - if verbose: - utils.log_and_message( - f"No file found at {search_path.as_posix()}" - ) - return [], [] - - all_folder_names, all_filenames = search_filesystem_path_for_folders( - search_path / search_prefix, return_full_path + rclone_config_name = cfg.get_rclone_config_name( + cfg["connection_method"] ) + + all_folder_names, all_filenames = search_local_or_remote( + search_path, + search_prefix, + rclone_config_name, + return_full_path, + ) + return all_folder_names, all_filenames -def search_gdrive_or_aws_for_folders( +def search_local_or_remote( search_path: Path, search_prefix: str, - cfg: Configs, + rclone_config_name: str | None, return_full_path: bool = False, ) -> Tuple[List[Any], List[Any]]: """Search for files and folders in central path using `rclone lsjson` command. This command lists all the files and folders in the central path in a json format. The json contains file/folder info about each file/folder like name, type, etc. + + Parameters + ---------- + search_path + The path to search (relative to the local or remote drive). For example, + for "local_filesystem" this is the path on the local machine. For "ssh", this + is the path on the machine that has been connected to. + search_prefix + The search string e.g. "sub-*". + rclone_config_name + Name of the rclone config for the remote (not set for local). `rclone config` + can be used in the terminal to see how rclone has stored these. In datashuttle, + these are managed by `Configs`. + return_full_path + If `True`, return the full filepath, otherwise return only the folder/file name. + """ + config_prefix = "" if not rclone_config_name else f"{rclone_config_name}:" + output = rclone.call_rclone( - "lsjson " - f"{cfg.get_rclone_config_name()}:{search_path.as_posix()} " - f'--include "{search_prefix}"', + f'lsjson {config_prefix}"{search_path.as_posix()}"', pipe_std=True, ) @@ -657,73 +663,26 @@ def search_gdrive_or_aws_for_folders( if output.returncode != 0: utils.log_and_message( - f"Error searching files at {search_path.as_posix()} \n {output.stderr.decode('utf-8') if output.stderr else ''}" + f"Error searching files at {search_path.as_posix()}\n" + f"{output.stderr.decode('utf-8') if output.stderr else ''}" ) return all_folder_names, all_filenames files_and_folders = json.loads(output.stdout) - try: - for file_or_folder in files_and_folders: - name = file_or_folder["Name"] - is_dir = file_or_folder.get("IsDir", False) - - to_append = ( - (search_path / name).as_posix() if return_full_path else name - ) - - if is_dir: - all_folder_names.append(to_append) - else: - all_filenames.append(to_append) - - except Exception: - utils.log_and_message( - f"Error searching files at {search_path.as_posix()}" - ) - - return all_folder_names, all_filenames - - -# Actual function implementation -def search_filesystem_path_for_folders( - search_path_with_prefix: Path, return_full_path: bool = False -) -> Tuple[List[Path | str], List[Path | str]]: - r"""Search a folder through the local filesystem. - - Use glob to search the full search path (including prefix) with glob. - Files are filtered out of results, returning folders only. - - Parameters - ---------- - search_path_with_prefix - Path to search along with search prefix e.g. "C:\drive\project\sub-*" - - return_full_path - If `True` returns the path to the discovered folder or file, - otherwise just the name. - - Returns - ------- - Discovered folders (`all_folder_names`) and files (`all_filenames`). + for file_or_folder in files_and_folders: + name = file_or_folder["Name"] - """ - all_folder_names = [] - all_filenames = [] + if not fnmatch.fnmatch(name, search_prefix): + continue - all_files_and_folders = list(glob.glob(search_path_with_prefix.as_posix())) - sorter_files_and_folders = sorted(all_files_and_folders) + is_dir = file_or_folder.get("IsDir", False) - for file_or_folder_str in sorter_files_and_folders: - file_or_folder = Path(file_or_folder_str) + to_append = search_path / name if return_full_path else name - if file_or_folder.is_dir(): - all_folder_names.append( - file_or_folder if return_full_path else file_or_folder.name - ) + if is_dir: + all_folder_names.append(to_append) else: - all_filenames.append( - file_or_folder if return_full_path else file_or_folder.name - ) + all_filenames.append(to_append) return all_folder_names, all_filenames diff --git a/datashuttle/utils/formatting.py b/datashuttle/utils/formatting.py index 625174218..0c77cddcf 100644 --- a/datashuttle/utils/formatting.py +++ b/datashuttle/utils/formatting.py @@ -64,7 +64,13 @@ def check_and_format_names( names_to_format, reserved_keywords = [], [] for name in names: if name in canonical_reserved_keywords() or tags("*") in name: - reserved_keywords.append(name) + if tags("to") in name: + # handle an edge case where use searches with both tags + reserved_keywords += update_names_with_range_to_flag( + [name], prefix + ) + else: + reserved_keywords.append(name) else: names_to_format.append(name) diff --git a/datashuttle/utils/gdrive.py b/datashuttle/utils/gdrive.py index d02479463..22c837707 100644 --- a/datashuttle/utils/gdrive.py +++ b/datashuttle/utils/gdrive.py @@ -1,9 +1,78 @@ -from datashuttle.utils import utils +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from datashuttle.configs.config_class import Configs + +from datashuttle.utils import rclone, utils # ----------------------------------------------------------------------------- # Helper Functions # ----------------------------------------------------------------------------- +# These functions are used by both API and TUI for setting up connections to google drive. + + +def preliminary_for_setup_without_browser( + cfg: Configs, + gdrive_client_secret: str | None, + rclone_config_name: str, + log: bool = True, +) -> str: + """Prepare rclone configuration for Google Drive without using a browser. + + This function prepares the rclone configuration for Google Drive without using a browser. + + The `config_is_local=false` flag tells rclone that the configuration process is being run + on a headless machine which does not have access to a browser. + + The `--non-interactive` flag is used to control rclone's behaviour while running it through + external applications. An `rclone config create` command would assume default values for config + variables in an interactive mode. If the `--non-interactive` flag is provided and rclone needs + the user to input some detail, a JSON blob will be returned with the question in it. For this + particular setup, rclone outputs a command for user to run on a machine with a browser. + + This function runs `rclone config create` with the user credentials and returns the rclone's output info. + This output info is presented to the user while asking for a `config_token`. + + Next, the user will run rclone's given command, authenticate with google drive and input the + config token given by rclone for datashuttle to proceed with the setup. + """ + client_id_key_value = ( + f"client_id {cfg['gdrive_client_id']} " + if cfg["gdrive_client_id"] + else " " + ) + client_secret_key_value = ( + f"client_secret {gdrive_client_secret} " + if gdrive_client_secret + else "" + ) + output = rclone.call_rclone( + f"config create " + f"{rclone_config_name} " + f"drive " + f"{client_id_key_value}" + f"{client_secret_key_value}" + f"scope drive " + f"root_folder_id {cfg['gdrive_root_folder_id']} " + f"config_is_local=false " + f"--non-interactive", + pipe_std=True, + ) + + # Extracting rclone's message from the json + output_json = json.loads(output.stdout) + message = output_json["Option"]["Help"] + + if log: + utils.log(message) + + return message + + # ----------------------------------------------------------------------------- # Python API # ----------------------------------------------------------------------------- @@ -27,14 +96,24 @@ def ask_user_for_browser(log: bool = True) -> bool: return answer -def prompt_and_get_service_account_filepath(log: bool = True): - """Get service account filepath from user.""" - message = "Please enter your service account file path: " - input_ = utils.get_user_input(message).strip() +def prompt_and_get_config_token( + cfg: Configs, + gdrive_client_secret: str | None, + rclone_config_name: str, + log: bool = True, +) -> str: + """Prompt the user for rclone config token. - if log: - utils.log(message) - utils.log(f"Service account file at: {input_}") + This function presents the rclone's output/message to ask the user to run a command, authenticate + with google drive and input the `config_token` generated by rclone. The `config_token` is + then used to complete rclone's config setup for google drive. + """ + message = preliminary_for_setup_without_browser( + cfg, gdrive_client_secret, rclone_config_name, log=log + ) + input_ = utils.get_user_input( + message + "\nEnter the output here: " + ).strip() return input_ diff --git a/datashuttle/utils/getters.py b/datashuttle/utils/getters.py index 33ee68f86..831a1a609 100644 --- a/datashuttle/utils/getters.py +++ b/datashuttle/utils/getters.py @@ -293,9 +293,9 @@ def get_existing_project_paths() -> List[Path]: """ datashuttle_path = canonical_folders.get_datashuttle_path() - all_folders, _ = folders.search_filesystem_path_for_folders( - datashuttle_path / "*" - ) + all_folders = [ + path_ for path_ in datashuttle_path.glob("*") if path_.is_dir() + ] existing_project_paths = [] for folder_name in all_folders: diff --git a/datashuttle/utils/rclone.py b/datashuttle/utils/rclone.py index 59ed05ee8..2f3979146 100644 --- a/datashuttle/utils/rclone.py +++ b/datashuttle/utils/rclone.py @@ -215,7 +215,7 @@ def setup_rclone_config_for_gdrive( cfg: Configs, rclone_config_name: str, gdrive_client_secret: str | None, - service_account_filepath: Optional[str] = None, + config_token: Optional[str] = None, ) -> subprocess.Popen: """Set up rclone config for connections to Google Drive. @@ -239,8 +239,8 @@ def setup_rclone_config_for_gdrive( gdrive_client_secret Google Drive client secret, mandatory when using a Google Drive client. - service_account_filepath : path to service account file path for connection - without browser + config_token : a token to setup rclone config without opening a browser, + needed if the user's machine doesn't have access to a browser. """ client_id_key_value = ( @@ -254,10 +254,10 @@ def setup_rclone_config_for_gdrive( else "" ) - service_account_filepath_arg = ( - "" - if service_account_filepath is None - else f"service_account_file {service_account_filepath}" + extra_args = ( + f"config_is_local=false config_token={config_token}" + if config_token + else "" ) process = call_rclone_with_popen( @@ -268,7 +268,7 @@ def setup_rclone_config_for_gdrive( f"{client_secret_key_value}" f"scope drive " f"root_folder_id {cfg['gdrive_root_folder_id']} " - f"{service_account_filepath_arg}", + f"{extra_args}", ) return process diff --git a/datashuttle/utils/ssh.py b/datashuttle/utils/ssh.py index 40f43ef53..074ca731c 100644 --- a/datashuttle/utils/ssh.py +++ b/datashuttle/utils/ssh.py @@ -5,10 +5,8 @@ if TYPE_CHECKING: from datashuttle.configs.config_class import Configs -import fnmatch -import stat from pathlib import Path -from typing import Any, List, Optional, Tuple +from typing import Optional import paramiko @@ -306,119 +304,3 @@ def verify_ssh_central_host( utils.log("Host not accepted. No connection made.") return success - - -# ----------------------------------------------------------------------------- -# Search over SSH -# ----------------------------------------------------------------------------- - - -def search_ssh_central_for_folders( - search_path: Path, - search_prefix: str, - cfg: Configs, - verbose: bool = True, - return_full_path: bool = False, -) -> Tuple[List[Any], List[Any]]: - """Search for the search prefix in the search path over SSH. - - Parameters - ---------- - search_path - Path to search for folders in. - - search_prefix - Search prefix for folder names e.g. "sub-*". - - cfg - See connect_client_with_logging(). - - verbose - If `True`, if a search folder cannot be found, a message - will be printed with the un-found path. - - return_full_path - include the search_path in the returned paths - - Returns - ------- - Discovered folders (`all_folder_names`) and files (`all_filenames`). - - """ - client: paramiko.SSHClient - with paramiko.SSHClient() as client: - connect_client_with_logging( - client, cfg, message_on_sucessful_connection=verbose - ) - - sftp = client.open_sftp() - - all_folder_names, all_filenames = get_list_of_folder_names_over_sftp( - sftp, - search_path, - search_prefix, - verbose, - return_full_path, - ) - - return all_folder_names, all_filenames - - -def get_list_of_folder_names_over_sftp( - sftp: paramiko.sftp_client.SFTPClient, - search_path: Path, - search_prefix: str, - verbose: bool = True, - return_full_path: bool = False, -) -> Tuple[List[Any], List[Any]]: - """Use paramiko's sftp to search a path over ssh for folders. - - Return the folder names. - - Parameters - ---------- - sftp - Connected paramiko stfp object - (see search_ssh_central_for_folders()). - - search_path - Path to search for folders in. - - search_prefix - Prefix (can include wildcards) - to search folder names. - - verbose - If `True`, if a search folder cannot be found, a message - will be printed with the un-found path. - - return_full_path - include the search_path in the returned paths. - - Returns - ------- - Discovered folders (`all_folder_names`) and files (`all_filenames`). - - """ - all_folder_names = [] - all_filenames = [] - try: - for file_or_folder in sftp.listdir_attr(search_path.as_posix()): - if file_or_folder.st_mode is not None and fnmatch.fnmatch( - file_or_folder.filename, search_prefix - ): - to_append = ( - search_path / file_or_folder.filename - if return_full_path - else file_or_folder.filename - ) - if stat.S_ISDIR(file_or_folder.st_mode): - all_folder_names.append(to_append) - else: - all_filenames.append(to_append) - - except FileNotFoundError: - if verbose: - utils.log_and_message(f"No file found at {search_path.as_posix()}") - - return all_folder_names, all_filenames diff --git a/tests/tests_transfers/local_filesystem/test_transfer_special_arguments.py b/tests/tests_transfers/local_filesystem/test_transfer_special_arguments.py index 71ca45d9e..4f01a8da4 100644 --- a/tests/tests_transfers/local_filesystem/test_transfer_special_arguments.py +++ b/tests/tests_transfers/local_filesystem/test_transfer_special_arguments.py @@ -1,5 +1,6 @@ """ """ +import fnmatch import shutil from pathlib import Path @@ -54,17 +55,138 @@ def test_combinations_filesystem_transfer( upload_or_download, ): """ - Test many combinations of possible file transfer commands. The - entire test project is created in the original `local_path` + Test many combinations of possible file transfer commands. + """ + pathtable, project = pathtable_and_project + + paths_to_transferred_files = self.perform_transfer( + project, upload_or_download, sub_names, ses_names, datatype + ) + + expected_transferred_paths = self.get_expected_transferred_paths( + pathtable, sub_names, ses_names, datatype + ) + + assert sorted(paths_to_transferred_files) == sorted( + expected_transferred_paths + ) + + # Test Wildcards + # ---------------------------------------------------------------------------------- + # It is very difficult to test wildcards using the original machinery + # for testing keywords such as "all", "all_sub" etc as used in test_combinations_filesystem_transfer(). + # Therefore, test a few specific cases here. + + @pytest.mark.parametrize("upload_or_download", ["upload", "download"]) + def test_local_filesystem_wildcards_1( + self, pathtable_and_project, upload_or_download + ): + """Test a single custom transfer that combines different special keywords.""" + pathtable, project = pathtable_and_project + + sub_names = ["@*@date@*@"] + ses_names = ["all_ses"] + datatype = ["funcimg"] + + paths_to_transferred_files = self.perform_transfer( + project, upload_or_download, sub_names, ses_names, datatype + ) + + pathtable = pathtable[ + pathtable["parent_sub"] + .fillna("") + .apply(lambda x: fnmatch.fnmatch(x, "*date*")) + ] + + pathtable = pathtable[ + pathtable["parent_datatype"].apply(lambda x: x == "funcimg") + ] + + expected_transferred_paths = pathtable["path"] + + assert sorted(paths_to_transferred_files) == sorted( + expected_transferred_paths + ) + + @pytest.mark.parametrize("upload_or_download", ["upload", "download"]) + def test_local_filesystem_wildcards_2( + self, pathtable_and_project, upload_or_download + ): + """Test a single custom transfer that combines different special keywords.""" + pathtable, project = pathtable_and_project + + sub_names = ["all_sub"] + ses_names = ["ses-003@*@"] + datatype = ["all_non_datatype"] + + paths_to_transferred_files = self.perform_transfer( + project, upload_or_download, sub_names, ses_names, datatype + ) + + pathtable = pathtable[ + pathtable["parent_ses"] + .fillna("") + .apply(lambda x: fnmatch.fnmatch(x, "ses-003*")) + ] + + pathtable = pathtable[ + pathtable["parent_datatype"].apply(lambda x: x is None) + ] + + expected_transferred_paths = pathtable["path"] + + assert sorted(paths_to_transferred_files) == sorted( + expected_transferred_paths + ) + + @pytest.mark.parametrize("upload_or_download", ["upload", "download"]) + def test_local_filesystem_wildcards_3( + self, pathtable_and_project, upload_or_download + ): + """Test a single custom transfer that combines different special keywords.""" + pathtable, project = pathtable_and_project + + sub_names = ["sub-002@TO@003_@*@"] + ses_names = ["ses-001"] + datatype = ["all"] + + paths_to_transferred_files = self.perform_transfer( + project, upload_or_download, sub_names, ses_names, datatype + ) + + pathtable = pathtable[ + pathtable["parent_sub"] + .fillna("") + .apply( + lambda x: fnmatch.fnmatch(x, "sub-002*") + or fnmatch.fnmatch(x, "sub-003*") + ) + ] + + pathtable = pathtable[ + pathtable["parent_ses"] + .fillna("") + .apply(lambda x: fnmatch.fnmatch(x, "ses-001")) + ] + + expected_transferred_paths = pathtable["path"] + + assert sorted(paths_to_transferred_files) == sorted( + expected_transferred_paths + ) + + def perform_transfer( + self, project, upload_or_download, sub_names, ses_names, datatype + ): + """Transfer the data, swapping the paths to move a subset of + files from the already set up directory to a new directory + using upload or download. + + The entire test project is created in the original `local_path` and subset of it is uploaded and tested against. To test upload vs. download, the `local_path` and `central_path` locations are swapped. """ - pathtable, project = pathtable_and_project - - # Transfer the data, swapping the paths to move a subset of - # files from the already set up directory to a new directory - # using upload or download. transfer_function = test_utils.handle_upload_or_download( project, upload_or_download, @@ -81,10 +203,6 @@ def test_combinations_filesystem_transfer( project, swap_last_folder_only=False ) - expected_transferred_paths = self.get_expected_transferred_paths( - pathtable, sub_names, ses_names, datatype - ) - # Check what paths were actually moved # (through the local filesystem), and test path_to_search = ( @@ -100,12 +218,10 @@ def test_combinations_filesystem_transfer( paths_to_transferred_files ) - assert sorted(paths_to_transferred_files) == sorted( - expected_transferred_paths - ) - # Teardown here, because we have session scope. try: shutil.rmtree(self.central_from_local(project.cfg["local_path"])) except FileNotFoundError: pass + + return paths_to_transferred_files diff --git a/tests/tests_transfers/ssh/test_ssh_transfer.py b/tests/tests_transfers/ssh/test_ssh_transfer.py index f086462ef..e9dbd089e 100644 --- a/tests/tests_transfers/ssh/test_ssh_transfer.py +++ b/tests/tests_transfers/ssh/test_ssh_transfer.py @@ -1,3 +1,4 @@ +import fnmatch import platform import shutil @@ -66,14 +67,117 @@ def test_combinations_ssh_transfer( to a container. This is very slow, due to the rclone ssh transfer (which is performed twice in this test, once for upload, once for download), around 8 seconds per parameterization. + """ + pathtable, project = ssh_setup + + expected_transferred_paths = self.get_expected_transferred_paths( + pathtable, sub_names, ses_names, datatype + ) + + self.check_transfer_ssh( + project, sub_names, ses_names, datatype, expected_transferred_paths + ) + + # Test Wildcards + # ---------------------------------------------------------------------------------- + # It is very difficult to test wildcards using the original machinery + # for testing keywords such as "all", "all_sub" etc as used in test_combinations_ssh_transfer(). + # Therefore, test a few specific cases here by manually chopping down the pathtable based + # on the sub / ses /datatype names to test the expected paths. + + def test_ssh_wildcards_1(self, ssh_setup): + """Test a single custom transfer that combines different special keywords.""" + pathtable, project = ssh_setup + + sub_names = ["@*@date@*@"] + ses_names = ["all_ses"] + datatype = ["funcimg"] + + pathtable = pathtable[ + pathtable["parent_sub"] + .fillna("") + .apply(lambda x: fnmatch.fnmatch(x, "*date*")) + ] + + pathtable = pathtable[ + pathtable["parent_datatype"].apply(lambda x: x == "funcimg") + ] + + expected_transferred_paths = pathtable["path"] + + self.check_transfer_ssh( + project, sub_names, ses_names, datatype, expected_transferred_paths + ) + + def test_ssh_wildcards_2(self, ssh_setup): + """Test a single custom transfer that combines different special keywords.""" + pathtable, project = ssh_setup + + sub_names = ["all_sub"] + ses_names = ["ses-003@*@"] + datatype = ["all_non_datatype"] + + pathtable = pathtable[ + pathtable["parent_ses"] + .fillna("") + .apply(lambda x: fnmatch.fnmatch(x, "ses-003*")) + ] + + pathtable = pathtable[ + pathtable["parent_datatype"].apply(lambda x: x is None) + ] + + expected_transferred_paths = pathtable["path"] + + self.check_transfer_ssh( + project, sub_names, ses_names, datatype, expected_transferred_paths + ) + + def test_ssh_wildcards_3(self, ssh_setup): + """Test a single custom transfer that combines different special keywords.""" + pathtable, project = ssh_setup + + sub_names = ["sub-002@TO@003_@*@"] + ses_names = ["ses-001"] + datatype = ["all"] + + pathtable = pathtable[ + pathtable["parent_sub"] + .fillna("") + .apply( + lambda x: fnmatch.fnmatch(x, "sub-002*") + or fnmatch.fnmatch(x, "sub-003*") + ) + ] + + pathtable = pathtable[ + pathtable["parent_ses"] + .fillna("") + .apply(lambda x: fnmatch.fnmatch(x, "ses-001")) + ] + + expected_transferred_paths = pathtable["path"] + + self.check_transfer_ssh( + project, sub_names, ses_names, datatype, expected_transferred_paths + ) + + def check_transfer_ssh( + self, + project, + sub_names, + ses_names, + datatype, + expected_transferred_paths, + ): + """Transfer the data and check the transferred files match the + `expected_transferred_paths`. In test setup, the entire project is created in the `local_path` and is uploaded to `central_path`. So we only need to set up once per test, upload and download is to temporary folders and these temporary folders are cleaned at the end of each parameterization. """ - pathtable, project = ssh_setup - # Upload data from the setup local project to a temporary # central directory. true_central_path = project.cfg["central_path"] @@ -88,10 +192,6 @@ def test_combinations_ssh_transfer( "rawdata", sub_names, ses_names, datatype, init_log=False ) - expected_transferred_paths = self.get_expected_transferred_paths( - pathtable, sub_names, ses_names, datatype - ) - # Search the paths that were transferred and tidy them up, # then check against the paths that were expected to be transferred. transferred_files = ssh_test_utils.recursive_search_central(project)