From 569e24513f8ee969aedf6bdea3c8c3abfb2a1305 Mon Sep 17 00:00:00 2001 From: Bob Date: Fri, 21 Aug 2026 09:07:34 -0600 Subject: [PATCH 1/2] Deduplicate CLI subcommands, parser setup, and repository commands Signed-off-by: Bob --- lib/ramble/ramble/cmd/common/__init__.py | 63 +++++++++++ lib/ramble/ramble/cmd/data.py | 42 ++------ lib/ramble/ramble/cmd/deployment.py | 41 ++----- lib/ramble/ramble/cmd/repo.py | 129 +++++++++++------------ lib/ramble/ramble/cmd/workspace.py | 123 ++++----------------- lib/ramble/ramble/test/cmd/common.py | 77 ++++++++++++++ 6 files changed, 239 insertions(+), 236 deletions(-) create mode 100644 lib/ramble/ramble/test/cmd/common.py diff --git a/lib/ramble/ramble/cmd/common/__init__.py b/lib/ramble/ramble/cmd/common/__init__.py index d0ab9f053d..45d9d2f14d 100644 --- a/lib/ramble/ramble/cmd/common/__init__.py +++ b/lib/ramble/ramble/cmd/common/__init__.py @@ -6,6 +6,8 @@ # option. This file may not be copied, modified, or distributed # except according to those terms. +import builtins + import ramble.paths import ramble.util.colors as color from ramble.util.logger import logger @@ -50,3 +52,64 @@ def shell_init_instructions(cmd, equivalent): msg += [""] logger.error(*msg) + + +def sanitize_arg_name(base_name: str) -> str: + """Format argument/command names by converting hyphens to underscores.""" + return base_name.replace("-", "_") + + +def setup_subcommands_from_prefix( + subparser, + dest: str, + subcommands: builtins.list, + prefix: str, + globals_dict: dict, + subcommand_functions: dict, + inject_dry_run: bool = False, +): + """Set up subparsers and map functions for a prefixed list of subcommands. + + Args: + subparser: Parent ArgumentParser or subparser + dest (str): Destination attribute name for selected subcommand + subcommands (list): List of subcommand names or (name, *aliases) tuples + prefix (str): Prefix used in function names (e.g. 'workspace', 'data', 'deployment') + globals_dict (dict): The globals() dictionary of the calling module + subcommand_functions (dict): Dictionary mapping command/alias name -> function + inject_dry_run (bool): Whether to automatically add --dry-run to subparsers if not present + """ + sp = subparser.add_subparsers(metavar="SUBCOMMAND", dest=dest) + + for cmd_entry in subcommands: + if isinstance(cmd_entry, (builtins.list, builtins.tuple)): + name, aliases = cmd_entry[0], builtins.list(cmd_entry[1:]) + else: + name = cmd_entry + aliases = [] + + # add commands to subcommands dict + function_name = sanitize_arg_name(f"{prefix}_{name}") + function = globals_dict[function_name] + for alias in [name] + aliases: + subcommand_functions[alias] = function + + # make a subparser and run the command's setup function on it + setup_parser_cmd_name = sanitize_arg_name(f"{prefix}_{name}_setup_parser") + setup_parser_cmd = globals_dict[setup_parser_cmd_name] + + subsubparser = sp.add_parser( + name, + aliases=aliases, + help=setup_parser_cmd.__doc__, + description=setup_parser_cmd.__doc__, + ) + setup_parser_cmd(subsubparser) + + if inject_dry_run and "--dry-run" not in subsubparser._option_string_actions: + subsubparser.add_argument( + "--dry-run", + dest="dry_run", + action="store_true", + help=f"perform a dry run of the {name} command", + ) diff --git a/lib/ramble/ramble/cmd/data.py b/lib/ramble/ramble/cmd/data.py index aebe4a043d..3e302ef955 100644 --- a/lib/ramble/ramble/cmd/data.py +++ b/lib/ramble/ramble/cmd/data.py @@ -8,6 +8,8 @@ from typing import Callable, Dict +import ramble.cmd.common + description = "manage data and databases" section = "data" level = "short" @@ -60,39 +62,15 @@ def data_create_db(args): subcommand_functions: Dict[str, Callable] = {} -def sanitize_arg_name(base_name): - """Allow function names to be remapped (eg `-` to `_`)""" - formatted_name = base_name.replace("-", "_") - return formatted_name - - def setup_parser(subparser): - sp = subparser.add_subparsers(metavar="SUBCOMMAND", dest="data_command") - - for name in subcommands: - if isinstance(name, (list, tuple)): - name, aliases = name[0], name[1:] - else: - aliases = [] - - # add commands to subcommands dict - function_name = sanitize_arg_name(f"data_{name}") - - function = globals()[function_name] - for alias in [name] + aliases: - subcommand_functions[alias] = function - - # make a subparser and run the command's setup function on it - setup_parser_cmd_name = sanitize_arg_name(f"data_{name}_setup_parser") - setup_parser_cmd = globals()[setup_parser_cmd_name] - - subsubparser = sp.add_parser( - name, - aliases=aliases, - help=setup_parser_cmd.__doc__, - description=setup_parser_cmd.__doc__, - ) - setup_parser_cmd(subsubparser) + ramble.cmd.common.setup_subcommands_from_prefix( + subparser=subparser, + dest="data_command", + subcommands=subcommands, + prefix="data", + globals_dict=globals(), + subcommand_functions=subcommand_functions, + ) def data(parser, args, unknown_args=None): diff --git a/lib/ramble/ramble/cmd/deployment.py b/lib/ramble/ramble/cmd/deployment.py index 27895a3aed..6e43806425 100644 --- a/lib/ramble/ramble/cmd/deployment.py +++ b/lib/ramble/ramble/cmd/deployment.py @@ -12,6 +12,7 @@ import llnl.util.filesystem as fs import ramble.cmd +import ramble.cmd.common import ramble.config import ramble.fetch_strategy import ramble.filters @@ -195,39 +196,15 @@ def deployment_run_pipeline(args, pipeline): subcommand_functions: Dict[str, Callable] = {} -def sanitize_arg_name(base_name): - """Allow function names to be remapped (eg `-` to `_`)""" - formatted_name = base_name.replace("-", "_") - return formatted_name - - def setup_parser(subparser): - sp = subparser.add_subparsers(metavar="SUBCOMMAND", dest="deployment_command") - - for name in subcommands: - if isinstance(name, (list, tuple)): - name, aliases = name[0], name[1:] - else: - aliases = [] - - # add commands to subcommands dict - function_name = sanitize_arg_name(f"deployment_{name}") - - function = globals()[function_name] - for alias in [name] + aliases: - subcommand_functions[alias] = function - - # make a subparser and run the command's setup function on it - setup_parser_cmd_name = sanitize_arg_name(f"deployment_{name}_setup_parser") - setup_parser_cmd = globals()[setup_parser_cmd_name] - - subsubparser = sp.add_parser( - name, - aliases=aliases, - help=setup_parser_cmd.__doc__, - description=setup_parser_cmd.__doc__, - ) - setup_parser_cmd(subsubparser) + ramble.cmd.common.setup_subcommands_from_prefix( + subparser=subparser, + dest="deployment_command", + subcommands=subcommands, + prefix="deployment", + globals_dict=globals(), + subcommand_functions=subcommand_functions, + ) def deployment(parser, args): diff --git a/lib/ramble/ramble/cmd/repo.py b/lib/ramble/ramble/cmd/repo.py index e0c4e03bc2..0b9128bea6 100644 --- a/lib/ramble/ramble/cmd/repo.py +++ b/lib/ramble/ramble/cmd/repo.py @@ -121,6 +121,63 @@ def repo_create(args): ) +def _add_repo_to_config(obj_type, canon_path, original_path, namespace, scope): + """Add a canonical repository path for an object type to configuration + if not already present.""" + type_def = ramble.repository.type_definitions[obj_type] + repos = ramble.config.get(type_def["config_section"], scope=scope) or [] + + is_present = any( + ramble.util.path.canonicalize_path(r_path) == canon_path or r_path == original_path + for r_path in repos + ) + + if is_present: + logger.warn( + f"{obj_type.name} repository is already registered with Ramble: {original_path}" + ) + else: + repos.insert(0, canon_path) + ramble.config.set(type_def["config_section"], repos, scope) + logger.msg(f"Added {obj_type.name} repo with namespace '{namespace}'.") + + +def _remove_repo_from_config(obj_type, namespace_or_path, scope): + """Remove a repository from a specific object type and scope by path or namespace.""" + type_def = ramble.repository.type_definitions[obj_type] + repos = ramble.config.get(type_def["config_section"], scope=scope) + if not repos: + return False + + canon_path = ramble.util.path.canonicalize_path(namespace_or_path) + normalized_path_to_remove = os.path.normcase(os.path.normpath(canon_path)) + + for repo_path in repos: + repo_canon_path = ramble.util.path.canonicalize_path(repo_path) + normalized_repo_path = os.path.normcase(os.path.normpath(repo_canon_path)) + if normalized_path_to_remove == normalized_repo_path: + repos.remove(repo_path) + ramble.config.set(type_def["config_section"], repos, scope) + logger.msg(f"Removed {obj_type.name} repository {repo_path} from scope '{scope}'.") + return True + + for path in list(repos): + try: + repo = ramble.repository.Repo(path, obj_type) + if repo.namespace == namespace_or_path: + repos.remove(path) + ramble.config.set(type_def["config_section"], repos, scope) + logger.msg( + f"Removed {obj_type.name} repository {repo.root} " + f"with namespace '{repo.namespace}' from scope '{scope}'." + ) + return True + except ramble.repository.RepoError: + continue + + return False + + def repo_add(args): """Add a repository to Ramble's configuration.""" path = args.path @@ -199,27 +256,9 @@ def repo_add(args): # Now add the canonical path for all object types for obj_type in ramble.repository.ObjectTypes: - type_def = ramble.repository.type_definitions[obj_type] - repos = ramble.config.get(type_def["config_section"], scope=args.scope) or [] - - # Check if canonical path is already present in the list of repos for this type - is_present = False - for r_path in repos: - if ramble.util.path.canonicalize_path(r_path) == canon_path: - is_present = True - break - - if is_present: - logger.warn( - f"{obj_type.name} repository is already registered with Ramble: {path}" - ) - else: - repos.insert(0, canon_path) - ramble.config.set(type_def["config_section"], repos, args.scope) - logger.msg(f"Added {obj_type.name} repo with namespace '{repo_namespace}'.") + _add_repo_to_config(obj_type, canon_path, path, repo_namespace, args.scope) else: # This is the original logic for a specific type obj_type = ramble.repository.ObjectTypes[args.type] - type_def = ramble.repository.type_definitions[obj_type] allow_partial = False # For specific type, we don't allow partial # Make sure it's actually a ramble repository by constructing it. @@ -241,14 +280,7 @@ def repo_add(args): f"The given path {path} is not a valid repo for type {obj_type.name}" ) - repos = ramble.config.get(type_def["config_section"], scope=args.scope) or [] - - if repo.root in repos or path in repos: - logger.warn(f"{obj_type.name} repository is already registered with Ramble: {path}") - else: - repos.insert(0, canon_path) - ramble.config.set(type_def["config_section"], repos, args.scope) - logger.msg(f"Added {obj_type.name} repo with namespace '{repo.namespace}'.") + _add_repo_to_config(obj_type, canon_path, path, repo.namespace, args.scope) def repo_remove(args): @@ -267,47 +299,8 @@ def repo_remove(args): repo_removed = False for scope in scopes_to_check: for obj_type in obj_types: - type_def = ramble.repository.type_definitions[obj_type] - repos = ramble.config.get(type_def["config_section"], scope=scope) - if not repos: - continue - - namespace_or_path = args.namespace_or_path - - canon_path = ramble.util.path.canonicalize_path(namespace_or_path) - normalized_path_to_remove = os.path.normcase(os.path.normpath(canon_path)) - - path_found_and_removed = False - for repo_path in repos: - repo_canon_path = ramble.util.path.canonicalize_path(repo_path) - normalized_repo_path = os.path.normcase(os.path.normpath(repo_canon_path)) - if normalized_path_to_remove == normalized_repo_path: - repos.remove(repo_path) - ramble.config.set(type_def["config_section"], repos, scope) - logger.msg( - f"Removed {obj_type.name} repository {repo_path} from scope '{scope}'." - ) - repo_removed = True - path_found_and_removed = True - break # move to next obj_type - - if path_found_and_removed: - continue - - for path in list(repos): - try: - repo = ramble.repository.Repo(path, obj_type) - if repo.namespace == namespace_or_path: - repos.remove(path) - ramble.config.set(type_def["config_section"], repos, scope) - logger.msg( - f"Removed {obj_type.name} repository {repo.root} " - f"with namespace '{repo.namespace}' from scope '{scope}'." - ) - repo_removed = True - break - except ramble.repository.RepoError: - continue + if _remove_repo_from_config(obj_type, args.namespace_or_path, scope): + repo_removed = True if repo_removed and not args.scope: break diff --git a/lib/ramble/ramble/cmd/workspace.py b/lib/ramble/ramble/cmd/workspace.py index 21f7b8636d..3b492874c2 100644 --- a/lib/ramble/ramble/cmd/workspace.py +++ b/lib/ramble/ramble/cmd/workspace.py @@ -20,6 +20,7 @@ from llnl.util.tty.colify import colified, colify import ramble.cmd +import ramble.cmd.common import ramble.cmd.filter_groups import ramble.config import ramble.expander @@ -1965,43 +1966,7 @@ def workspace_manage_modifiers(args): def workspace_manage_filter_groups_setup_parser(subparser): """manage workspace filter groups""" - scopes_metavar = ramble.config.scopes_metavar - - subparser.add_argument( - "--scope", - choices=ramble.config.scopes_choices(include_workspace=True), - metavar=scopes_metavar, - default=None, - help="configuration scope to modify/list", - ) - - actions = subparser.add_subparsers(metavar="ACTION", dest="action") - - add_parser = actions.add_parser("add", help="add a filter group") - add_parser.add_argument("-n", "--name", required=True, help="name of filter group") - add_parser.add_argument( - "--where", - action="append", - help="inclusive filter expression. Can be specified multiple times.", - ) - add_parser.add_argument( - "--exclude-where", - dest="exclude_where", - action="append", - help="exclusive filter expression. Can be specified multiple times.", - ) - - remove_parser = actions.add_parser("remove", aliases=["rm"], help="remove a filter group") - remove_parser.add_argument("-n", "--name", required=True, help="name of filter group") - - list_parser = actions.add_parser("list", help="list defined filter groups") - list_parser.add_argument( - "-v", - "--verbose", - action="store_true", - help="show the filter group definition for each", - ) - actions.add_parser("blame", help="show defined filter groups with sources") + ramble.cmd.filter_groups.setup_parser(subparser) def workspace_manage_filter_groups(args): @@ -2118,48 +2083,16 @@ def workspace_experiment_logs(args): subcommand_functions: Dict[str, Callable] = {} -def sanitize_arg_name(base_name): - """Allow function names to be remapped (eg `-` to `_`)""" - formatted_name = base_name.replace("-", "_") - return formatted_name - - def setup_parser(subparser): - sp = subparser.add_subparsers(metavar="SUBCOMMAND", dest="workspace_command") - - for name in subcommands: - if isinstance(name, (list, tuple)): - name, aliases = name[0], name[1:] - else: - aliases = [] - - # add commands to subcommands dict - function_name = sanitize_arg_name(f"workspace_{name}") - - function = globals()[function_name] - for alias in [name] + aliases: - subcommand_functions[alias] = function - - # make a subparser and run the command's setup function on it - setup_parser_cmd_name = sanitize_arg_name(f"workspace_{name}_setup_parser") - setup_parser_cmd = globals()[setup_parser_cmd_name] - - subsubparser = sp.add_parser( - name, - aliases=aliases, - help=setup_parser_cmd.__doc__, - description=setup_parser_cmd.__doc__, - ) - setup_parser_cmd(subsubparser) - - # inject --dry-run into subcommands - if "--dry-run" not in subsubparser._option_string_actions: - subsubparser.add_argument( - "--dry-run", - dest="dry_run", - action="store_true", - help=f"perform a dry run of the {name} command", - ) + ramble.cmd.common.setup_subcommands_from_prefix( + subparser=subparser, + dest="workspace_command", + subcommands=subcommands, + prefix="workspace", + globals_dict=globals(), + subcommand_functions=subcommand_functions, + inject_dry_run=True, + ) def workspace(parser, args, unknown_args): @@ -2185,29 +2118,11 @@ def workspace_manage(args): def workspace_manage_setup_parser(subparser): """manage workspace definitions""" - sp = subparser.add_subparsers(metavar="SUBCOMMAND", dest="manage_command") - - for name in manage_commands: - if isinstance(name, (list, tuple)): - name, aliases = name[0], name[1:] - else: - aliases = [] - - # add commands to subcommands dict - function_name = sanitize_arg_name(f"workspace_manage_{name}") - - function = globals()[function_name] - for alias in [name] + aliases: - manage_subcommand_functions[alias] = function - - # make a subparser and run the command's setup function on it - setup_parser_cmd_name = sanitize_arg_name(f"workspace_manage_{name}_setup_parser") - setup_parser_cmd = globals()[setup_parser_cmd_name] - - subsubparser = sp.add_parser( - name, - aliases=aliases, - help=setup_parser_cmd.__doc__, - description=setup_parser_cmd.__doc__, - ) - setup_parser_cmd(subsubparser) + ramble.cmd.common.setup_subcommands_from_prefix( + subparser=subparser, + dest="manage_command", + subcommands=manage_commands, + prefix="workspace_manage", + globals_dict=globals(), + subcommand_functions=manage_subcommand_functions, + ) diff --git a/lib/ramble/ramble/test/cmd/common.py b/lib/ramble/ramble/test/cmd/common.py new file mode 100644 index 0000000000..167d2af27c --- /dev/null +++ b/lib/ramble/ramble/test/cmd/common.py @@ -0,0 +1,77 @@ +# Copyright 2022-2026 The Ramble Authors +# +# Licensed under the Apache License, Version 2.0 or the MIT license +# , at your +# option. This file may not be copied, modified, or distributed +# except according to those terms. + +import argparse + +import ramble.cmd.common + + +def test_sanitize_arg_name(): + assert ramble.cmd.common.sanitize_arg_name("foo-bar") == "foo_bar" + assert ramble.cmd.common.sanitize_arg_name("simple") == "simple" + assert ramble.cmd.common.sanitize_arg_name("a-b-c_d") == "a_b_c_d" + + +def test_setup_subcommands_from_prefix(): + def mock_cmd_foo_setup_parser(subparser): + """Foo command docstring""" + subparser.add_argument("--test-opt", action="store_true") + + def mock_cmd_foo(args): + return "foo executed" + + def mock_cmd_bar_baz_setup_parser(subparser): + """Bar Baz command docstring""" + subparser.add_argument("--bar-opt", type=str) + + def mock_cmd_bar_baz(args): + return "bar-baz executed" + + mock_globals = { + "mock_cmd_foo_setup_parser": mock_cmd_foo_setup_parser, + "mock_cmd_foo": mock_cmd_foo, + "mock_cmd_bar_baz_setup_parser": mock_cmd_bar_baz_setup_parser, + "mock_cmd_bar_baz": mock_cmd_bar_baz, + } + subcommand_functions = {} + subcommands = [ + "foo", + ("bar-baz", "bb", "barbaz"), + ] + + parser = argparse.ArgumentParser() + ramble.cmd.common.setup_subcommands_from_prefix( + subparser=parser, + dest="subcommand", + subcommands=subcommands, + prefix="mock_cmd", + globals_dict=mock_globals, + subcommand_functions=subcommand_functions, + inject_dry_run=True, + ) + + assert "foo" in subcommand_functions + assert "bar-baz" in subcommand_functions + assert "bb" in subcommand_functions + assert "barbaz" in subcommand_functions + assert subcommand_functions["bb"] == mock_cmd_bar_baz + + args = parser.parse_args(["foo", "--test-opt", "--dry-run"]) + assert args.subcommand == "foo" + assert args.test_opt is True + assert args.dry_run is True + + args_alias = parser.parse_args(["bb", "--bar-opt", "hello"]) + assert args_alias.subcommand == "bb" + assert args_alias.bar_opt == "hello" + + +def test_shell_init_instructions(caplog): + ramble.cmd.common.shell_init_instructions( + "workspace activate", "ramble workspace activate {sh_arg}" + ) From 516d4d65369cdb7ad6a723c42ad480de88bbe335 Mon Sep 17 00:00:00 2001 From: Bob Date: Fri, 21 Aug 2026 09:24:38 -0600 Subject: [PATCH 2/2] Apply gemini feedback to use built-in type hints Signed-off-by: Bob --- lib/ramble/ramble/cmd/common/__init__.py | 73 ++++------------------- lib/ramble/ramble/cmd/common/arguments.py | 69 ++++++++++++++++++++- 2 files changed, 78 insertions(+), 64 deletions(-) diff --git a/lib/ramble/ramble/cmd/common/__init__.py b/lib/ramble/ramble/cmd/common/__init__.py index 45d9d2f14d..0e5dd2c9a4 100644 --- a/lib/ramble/ramble/cmd/common/__init__.py +++ b/lib/ramble/ramble/cmd/common/__init__.py @@ -6,12 +6,20 @@ # option. This file may not be copied, modified, or distributed # except according to those terms. -import builtins - import ramble.paths import ramble.util.colors as color +from ramble.cmd.common.arguments import ( + sanitize_arg_name, + setup_subcommands_from_prefix, +) from ramble.util.logger import logger +__all__ = [ + "shell_init_instructions", + "sanitize_arg_name", + "setup_subcommands_from_prefix", +] + def shell_init_instructions(cmd, equivalent): """Print out instructions for users to initialize shell support. @@ -52,64 +60,3 @@ def shell_init_instructions(cmd, equivalent): msg += [""] logger.error(*msg) - - -def sanitize_arg_name(base_name: str) -> str: - """Format argument/command names by converting hyphens to underscores.""" - return base_name.replace("-", "_") - - -def setup_subcommands_from_prefix( - subparser, - dest: str, - subcommands: builtins.list, - prefix: str, - globals_dict: dict, - subcommand_functions: dict, - inject_dry_run: bool = False, -): - """Set up subparsers and map functions for a prefixed list of subcommands. - - Args: - subparser: Parent ArgumentParser or subparser - dest (str): Destination attribute name for selected subcommand - subcommands (list): List of subcommand names or (name, *aliases) tuples - prefix (str): Prefix used in function names (e.g. 'workspace', 'data', 'deployment') - globals_dict (dict): The globals() dictionary of the calling module - subcommand_functions (dict): Dictionary mapping command/alias name -> function - inject_dry_run (bool): Whether to automatically add --dry-run to subparsers if not present - """ - sp = subparser.add_subparsers(metavar="SUBCOMMAND", dest=dest) - - for cmd_entry in subcommands: - if isinstance(cmd_entry, (builtins.list, builtins.tuple)): - name, aliases = cmd_entry[0], builtins.list(cmd_entry[1:]) - else: - name = cmd_entry - aliases = [] - - # add commands to subcommands dict - function_name = sanitize_arg_name(f"{prefix}_{name}") - function = globals_dict[function_name] - for alias in [name] + aliases: - subcommand_functions[alias] = function - - # make a subparser and run the command's setup function on it - setup_parser_cmd_name = sanitize_arg_name(f"{prefix}_{name}_setup_parser") - setup_parser_cmd = globals_dict[setup_parser_cmd_name] - - subsubparser = sp.add_parser( - name, - aliases=aliases, - help=setup_parser_cmd.__doc__, - description=setup_parser_cmd.__doc__, - ) - setup_parser_cmd(subsubparser) - - if inject_dry_run and "--dry-run" not in subsubparser._option_string_actions: - subsubparser.add_argument( - "--dry-run", - dest="dry_run", - action="store_true", - help=f"perform a dry run of the {name} command", - ) diff --git a/lib/ramble/ramble/cmd/common/arguments.py b/lib/ramble/ramble/cmd/common/arguments.py index 1a8485d5a9..0b807c8265 100644 --- a/lib/ramble/ramble/cmd/common/arguments.py +++ b/lib/ramble/ramble/cmd/common/arguments.py @@ -15,7 +15,13 @@ from spack.util.pattern import Args -__all__ = ["add_common_arguments", "allows_unknown_args", "validate_unknown_args"] +__all__ = [ + "add_common_arguments", + "allows_unknown_args", + "validate_unknown_args", + "sanitize_arg_name", + "setup_subcommands_from_prefix", +] #: dictionary of argument-generating functions, keyed by name _arguments: Dict[str, Callable[[], Args]] = {} @@ -238,3 +244,64 @@ def exclude_filter_group(): help="Exclude experiments matching a logical expression of filter groups", required=False, ) + + +def sanitize_arg_name(base_name: str) -> str: + """Format argument/command names by converting hyphens to underscores.""" + return base_name.replace("-", "_") + + +def setup_subcommands_from_prefix( + subparser, + dest: str, + subcommands: list, + prefix: str, + globals_dict: dict, + subcommand_functions: dict, + inject_dry_run: bool = False, +): + """Set up subparsers and map functions for a prefixed list of subcommands. + + Args: + subparser: Parent ArgumentParser or subparser + dest (str): Destination attribute name for selected subcommand + subcommands (list): List of subcommand names or (name, *aliases) tuples + prefix (str): Prefix used in function names (e.g. 'workspace', 'data', 'deployment') + globals_dict (dict): The globals() dictionary of the calling module + subcommand_functions (dict): Dictionary mapping command/alias name -> function + inject_dry_run (bool): Whether to automatically add --dry-run to subparsers if not present + """ + sp = subparser.add_subparsers(metavar="SUBCOMMAND", dest=dest) + + for cmd_entry in subcommands: + if isinstance(cmd_entry, (list, tuple)): + name, aliases = cmd_entry[0], list(cmd_entry[1:]) + else: + name = cmd_entry + aliases = [] + + # add commands to subcommands dict + function_name = sanitize_arg_name(f"{prefix}_{name}") + function = globals_dict[function_name] + for alias in [name] + aliases: + subcommand_functions[alias] = function + + # make a subparser and run the command's setup function on it + setup_parser_cmd_name = sanitize_arg_name(f"{prefix}_{name}_setup_parser") + setup_parser_cmd = globals_dict[setup_parser_cmd_name] + + subsubparser = sp.add_parser( + name, + aliases=aliases, + help=setup_parser_cmd.__doc__, + description=setup_parser_cmd.__doc__, + ) + setup_parser_cmd(subsubparser) + + if inject_dry_run and "--dry-run" not in subsubparser._option_string_actions: + subsubparser.add_argument( + "--dry-run", + dest="dry_run", + action="store_true", + help=f"perform a dry run of the {name} command", + )