From 070b58b12d11206ff3908899a2564faf980f494c Mon Sep 17 00:00:00 2001 From: luandalmazo Date: Thu, 30 Jul 2026 09:51:37 -0300 Subject: [PATCH] add arg completion --- datamint/__main__.py | 54 ++++++++ datamint/client_cmd_tools/datamint_config.py | 129 +++++++++++++++--- datamint/client_cmd_tools/datamint_example.py | 20 ++- .../client_cmd_tools/datamint_inference.py | 20 ++- datamint/client_cmd_tools/datamint_train.py | 20 ++- datamint/client_cmd_tools/datamint_upload.py | 21 ++- pyproject.toml | 1 + 7 files changed, 241 insertions(+), 24 deletions(-) diff --git a/datamint/__main__.py b/datamint/__main__.py index 90ed54d4..b913f3ba 100644 --- a/datamint/__main__.py +++ b/datamint/__main__.py @@ -3,6 +3,7 @@ import argparse import importlib import importlib.metadata +import os import sys _COMMANDS: dict[str, str] = { @@ -14,6 +15,15 @@ "example": "datamint.client_cmd_tools.datamint_example", } +_COMMAND_HELP: dict[str, str] = { + "config": "Configure the API key, URL, and local cache", + "upload": "Upload DICOM files and other resources", + "init": "Generate starter scripts for a Datamint workflow", + "train": "Train a model on a Datamint project", + "inference": "Run local inference with a registered model", + "example": "Populate a project with an example dataset", +} + def _resolve_version() -> str: try: @@ -42,7 +52,51 @@ def _build_parser() -> argparse.ArgumentParser: return parser +def _completing_command_hint() -> str | None: + """Best-effort extraction of the subcommand token from COMP_LINE. """ + comp_line = os.environ.get("COMP_LINE", "") + tokens = comp_line.split() + if len(tokens) >= 2 and tokens[1] in _COMMANDS: + return tokens[1] + return None + + +def _autocomplete() -> None: + """Build a combined parser tree and hand it to argcomplete. """ + import argcomplete + + parser = argparse.ArgumentParser( + prog="datamint", + description="Datamint command-line interface.", + epilog=f"Available commands: {', '.join(_COMMANDS)}.", + ) + parser.add_argument("--version", action="version", version=f"%(prog)s {_resolve_version()}") + subparsers = parser.add_subparsers(dest="command", metavar="command") + + target = _completing_command_hint() + for name, module_path in _COMMANDS.items(): + if name != target: + subparsers.add_parser(name, help=_COMMAND_HELP[name]) + continue + try: + module = importlib.import_module(module_path) + except ImportError: + subparsers.add_parser(name, help=_COMMAND_HELP[name]) + continue + build_fn = getattr(module, "_build_parser", None) + if build_fn is not None: + build_fn(subparsers) + else: + subparsers.add_parser(name, help=_COMMAND_HELP[name]) + + argcomplete.autocomplete(parser) + + def main() -> None: + if os.environ.get("_ARGCOMPLETE"): + _autocomplete() + return + parser = _build_parser() argv = sys.argv[1:] diff --git a/datamint/client_cmd_tools/datamint_config.py b/datamint/client_cmd_tools/datamint_config.py index 200de763..f55894bc 100644 --- a/datamint/client_cmd_tools/datamint_config.py +++ b/datamint/client_cmd_tools/datamint_config.py @@ -720,6 +720,7 @@ def interactive_mode(): console.print(" [accent](4)[/accent] Clear all configuration settings") console.print(" [accent](5)[/accent] Test connection") console.print(" [accent](6)[/accent] Manage/Show local data...") + console.print(" [accent](7)[/accent] Set up shell tab-completion") console.print(" [accent](q)[/accent] Exit") choice = Prompt.ask("Enter your choice", console=console).lower().strip() @@ -735,6 +736,8 @@ def interactive_mode(): test_connection() elif choice == '6': interactive_dataset_cleaning() + elif choice == '7': + install_shell_completion() elif choice in ('q', 'exit', 'quit'): break else: @@ -745,21 +748,13 @@ def interactive_mode(): console.print("[success]👋 Goodbye![/success]") -def main(): - """Main entry point for the configuration tool.""" - global console - load_cmdline_logging_config() - console_handlers = [h for h in _USER_LOGGER.handlers if isinstance(h, ConsoleWrapperHandler)] - if console_handlers: - console = console_handlers[0].console +def _build_parser(subparsers: argparse._SubParsersAction | None = None) -> argparse.ArgumentParser: + """Build the argument parser. - if is_legacy_cli_invocation('config'): - console.print( - "[warning]'datamint-config' is deprecated and will be removed in a future " - "release. Use 'datamint config' instead.[/warning]" - ) - - parser = argparse.ArgumentParser( + When ``subparsers`` is given, the parser is registered as a ``config`` subparser + (used by ``datamint``'s combined completion tree) instead of a standalone parser. + """ + kwargs = dict( description='🔧 Datamint API Configuration Tool', epilog=""" Examples: @@ -775,11 +770,16 @@ def main(): datamint config --clean-local-data "Example Project" # Clean a legacy dataset folder datamint config --clean-all-local-data # Clean all local data groups - + datamint config --install-completion # Set up shell tab-completion + More Documentation: https://sonanceai.github.io/datamint-python-api/command_line_tools.html """, formatter_class=argparse.RawDescriptionHelpFormatter ) + if subparsers is not None: + parser = subparsers.add_parser('config', **kwargs) + else: + parser = argparse.ArgumentParser(**kwargs) parser.add_argument('--api-key', type=str, help='API key to set') parser.add_argument('--default-url', '--url', type=str, help='Default URL to set') parser.add_argument('-i', '--interactive', action='store_true', @@ -803,7 +803,101 @@ def main(): action='store_true', help='Clean all local cache namespaces and legacy dataset folders', ) + parser.add_argument( + '--install-completion', + dest='install_completion', + nargs='?', + const='auto', + choices=['auto', 'bash', 'fish'], + metavar='SHELL', + help='Set up shell tab-completion for the datamint CLI tools. ' + 'Auto-detects your shell from $SHELL if SHELL is omitted.', + ) + + return parser + + +_COMPLETION_EXECUTABLES = [ + 'datamint', 'datamint-upload', 'datamint-config', 'datamint-train', 'datamint-inference', +] + + +def _detect_shell() -> str | None: + """Best-effort shell detection from the $SHELL environment variable.""" + import os + shell_path = os.environ.get('SHELL', '') + name = Path(shell_path).name if shell_path else '' + return name if name in ('bash', 'fish', 'zsh') else None + + +def _completion_target_path(shell: str) -> Path: + """Return the conventional per-user completion file path for the given shell.""" + if shell == 'fish': + return Path.home() / '.config' / 'fish' / 'completions' / 'datamint.fish' + if shell == 'bash': + return Path.home() / '.local' / 'share' / 'bash-completion' / 'completions' / 'datamint' + raise ValueError(f"Unsupported shell for automatic install: {shell}") + + +def install_shell_completion(shell: str | None = None) -> None: + """Write shell tab-completion hooks for the datamint CLI entry points. """ + import argcomplete + + if shell is None or shell == 'auto': + detected = _detect_shell() + if detected is None: + console.print( + "[error]❌ Could not auto-detect your shell from $SHELL. " + "Pass it explicitly, e.g. --install-completion bash or --install-completion fish.[/error]" + ) + return + shell = detected + + if shell == 'zsh': + console.print("[warning]⚠️ Add these lines to your ~/.zshrc:[/warning]") + console.print(" autoload -U bashcompinit") + console.print(" bashcompinit") + for executable in _COMPLETION_EXECUTABLES: + console.print(f' eval "$(register-python-argcomplete {executable})"') + return + + if shell not in ('bash', 'fish'): + console.print( + f"[warning]⚠️ Automatic installation isn't supported for '{shell}' yet. " + "Add this line to your shell config instead:[/warning]" + ) + console.print(' eval "$(register-python-argcomplete datamint)"') + return + + script = argcomplete.shellcode(_COMPLETION_EXECUTABLES, shell=shell) + target = _completion_target_path(shell) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(script) + + console.print(f"[success]✅ Tab-completion installed at {target}[/success]") + console.print( + "[dim]Open a new terminal (or restart this one), then press Tab while typing a 'datamint' " + "command — it will suggest subcommands (config, upload, train, ...) and their flags.[/dim]" + ) + + +def main(): + """Main entry point for the configuration tool.""" + global console + load_cmdline_logging_config() + console_handlers = [h for h in _USER_LOGGER.handlers if isinstance(h, ConsoleWrapperHandler)] + if console_handlers: + console = console_handlers[0].console + if is_legacy_cli_invocation('config'): + console.print( + "[warning]'datamint-config' is deprecated and will be removed in a future " + "release. Use 'datamint config' instead.[/warning]" + ) + + parser = _build_parser() + import argcomplete + argcomplete.autocomplete(parser) args = parser.parse_args() config_updates: dict[str, str] = {} @@ -834,9 +928,12 @@ def main(): if args.clean_all_local_data: clean_all_datasets() + if args.install_completion is not None: + install_shell_completion(None if args.install_completion == 'auto' else args.install_completion) + no_arguments_provided = (args.api_key is None and args.default_url is None and not args.list_local_data and not args.clean_local_data and - not args.clean_all_local_data) + not args.clean_all_local_data and args.install_completion is None) if no_arguments_provided or args.interactive: interactive_mode() diff --git a/datamint/client_cmd_tools/datamint_example.py b/datamint/client_cmd_tools/datamint_example.py index a7952b8c..5dfafe55 100644 --- a/datamint/client_cmd_tools/datamint_example.py +++ b/datamint/client_cmd_tools/datamint_example.py @@ -18,8 +18,13 @@ _DATASETS = ('bccd', 'busi', 'synapse', 'fracatlas') -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( +def _build_parser(subparsers: argparse._SubParsersAction | None = None) -> argparse.ArgumentParser: + """Build the argument parser. + + When ``subparsers`` is given, the parser is registered as an ``example`` subparser + (used by ``datamint``'s combined completion tree) instead of a standalone parser. + """ + kwargs = dict( description='Populate a Datamint project with an example dataset.', epilog=""" Examples: @@ -32,11 +37,22 @@ def _parse_args() -> argparse.Namespace: """, formatter_class=argparse.RawDescriptionHelpFormatter, ) + if subparsers is not None: + parser = subparsers.add_parser('example', **kwargs) + else: + parser = argparse.ArgumentParser(**kwargs) parser.add_argument('dataset', choices=_DATASETS, help='Which example dataset to populate.') parser.add_argument('--project', type=str, default=None, help='Name of the project to create. Defaults to a dataset-specific name.') parser.add_argument('--verbose', action='store_true', default=False, help='Print debug messages.') + return parser + + +def _parse_args() -> argparse.Namespace: + parser = _build_parser() + import argcomplete + argcomplete.autocomplete(parser) return parser.parse_args() diff --git a/datamint/client_cmd_tools/datamint_inference.py b/datamint/client_cmd_tools/datamint_inference.py index de3b3439..fed56f4b 100644 --- a/datamint/client_cmd_tools/datamint_inference.py +++ b/datamint/client_cmd_tools/datamint_inference.py @@ -158,8 +158,13 @@ def _execute(args: argparse.Namespace, console: Console) -> int: return 0 -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( +def _build_parser(subparsers: argparse._SubParsersAction | None = None) -> argparse.ArgumentParser: + """Build the argument parser. + + When ``subparsers`` is given, the parser is registered as an ``inference`` subparser + (used by ``datamint``'s combined completion tree) instead of a standalone parser. + """ + kwargs = dict( description='Run local inference with a registered Datamint model against a local file.', epilog=""" Examples: @@ -176,6 +181,10 @@ def _parse_args() -> argparse.Namespace: """, formatter_class=argparse.RawDescriptionHelpFormatter, ) + if subparsers is not None: + parser = subparsers.add_parser('inference', **kwargs) + else: + parser = argparse.ArgumentParser(**kwargs) parser.add_argument('file', type=_is_valid_path_argparse, metavar='FILE', help='Path to the local file to run inference on.') parser.add_argument('--model-name', type=str, required=True, @@ -191,6 +200,13 @@ def _parse_args() -> argparse.Namespace: '(see datamint.utils.uncertainty). Off by default.') parser.add_argument('--verbose', action='store_true', default=False, help='Print debug messages.') + return parser + + +def _parse_args() -> argparse.Namespace: + parser = _build_parser() + import argcomplete + argcomplete.autocomplete(parser) return parser.parse_args() diff --git a/datamint/client_cmd_tools/datamint_train.py b/datamint/client_cmd_tools/datamint_train.py index b8c7a568..01d59190 100644 --- a/datamint/client_cmd_tools/datamint_train.py +++ b/datamint/client_cmd_tools/datamint_train.py @@ -351,8 +351,13 @@ def _execute(args: argparse.Namespace, api: Api, console: Console, *, show_plan: return 0 -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( +def _build_parser(subparsers: argparse._SubParsersAction | None = None) -> argparse.ArgumentParser: + """Build the argument parser. + + When ``subparsers`` is given, the parser is registered as a ``train`` subparser + (used by ``datamint``'s combined completion tree) instead of a standalone parser. + """ + kwargs = dict( description='Train a model on a Datamint project using a built-in one-line trainer.', epilog=""" Examples: @@ -367,6 +372,10 @@ def _parse_args() -> argparse.Namespace: """, formatter_class=argparse.RawDescriptionHelpFormatter, ) + if subparsers is not None: + parser = subparsers.add_parser('train', **kwargs) + else: + parser = argparse.ArgumentParser(**kwargs) parser.add_argument('--project', type=str, help='Project name to train on.') parser.add_argument('--model', type=str, choices=sorted(MODEL_REGISTRY), help='Model to train: ' + ', '.join(sorted(MODEL_REGISTRY)) + @@ -387,6 +396,13 @@ def _parse_args() -> argparse.Namespace: parser.add_argument('--interactive', action='store_true', help='Guided interactive wizard.') parser.add_argument('--verbose', action='store_true', default=False, help='Print debug messages.') + return parser + + +def _parse_args() -> argparse.Namespace: + parser = _build_parser() + import argcomplete + argcomplete.autocomplete(parser) args = parser.parse_args() if not args.interactive and not args.project: diff --git a/datamint/client_cmd_tools/datamint_upload.py b/datamint/client_cmd_tools/datamint_upload.py index d15c811b..653220dd 100644 --- a/datamint/client_cmd_tools/datamint_upload.py +++ b/datamint/client_cmd_tools/datamint_upload.py @@ -519,9 +519,18 @@ def _get_files_from_path(path: str | Path, raise -def _parse_args() -> tuple[Any, list[str], list[dict] | None, list[str] | None]: - parser = argparse.ArgumentParser( +def _build_parser(subparsers: argparse._SubParsersAction | None = None) -> argparse.ArgumentParser: + """Build the argument parser. + + When ``subparsers`` is given, the parser is registered as an ``upload`` subparser + (used by ``datamint``'s combined completion tree) instead of a standalone parser. + """ + kwargs = dict( description='DatamintAPI command line tool for uploading DICOM files and other resources') + if subparsers is not None: + parser = subparsers.add_parser('upload', **kwargs) + else: + parser = argparse.ArgumentParser(**kwargs) # Add positional argument for path parser.add_argument('path', nargs='?', type=_is_valid_path_argparse, metavar="PATH", @@ -591,6 +600,14 @@ def _parse_args() -> tuple[Any, list[str], list[dict] | None, list[str] | None]: 'By default, common non-medical file extensions are excluded unless --include-extensions is used.') parser.add_argument('--version', action='version', version=f'%(prog)s {datamint_version}') parser.add_argument('--verbose', action='store_true', help='Print debug messages', default=False) + + return parser + + +def _parse_args() -> tuple[Any, list[str], list[dict] | None, list[str] | None]: + parser = _build_parser() + import argcomplete + argcomplete.autocomplete(parser) args = parser.parse_args() # Handle path argument priority: positional takes precedence over --path flag diff --git a/pyproject.toml b/pyproject.toml index 59c605a5..962901c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ typing_extensions = ">=4.0.0" pydantic = ">=2.6.4" certifi = ">=2025.0.0" httpx = "*" +argcomplete = ">=3.0.0" cachetools = ">=6.0.0" segmentation-models-pytorch = ">=0.5.0" backports-strenum = { version = "*", python = "<3.11" }