diff --git a/shtab/__init__.py b/shtab/__init__.py index 2c34b7b..4e700af 100644 --- a/shtab/__init__.py +++ b/shtab/__init__.py @@ -141,6 +141,16 @@ def wordify(string): return string.replace("-", "_").replace(".", " ").replace(" ", "_") +def get_public_subcommands(sub): + """Get all the publicly-visible subcommands for a given subparser.""" + # NOTE: public subcommands have their primary name listed in the result of + # `_get_subactions()`. We use this to get the parser for each subcommand and + # compare all the choices (including aliases!) to the set of known-public + # parsers. + public_parsers = {id(sub.choices[i.dest]) for i in sub._get_subactions()} + return {k for k, v in sub.choices.items() if id(v) in public_parsers} + + def get_bash_commands(root_parser, root_prefix, choice_functions=None): """ Recursive subcommand parser traversal, returning lists of information on @@ -204,6 +214,9 @@ def recurse(parser, prefix): # choices (including subparsers & shtab `.complete` functions) log.debug("choices:{}:{}".format(prefix, sorted(positional.choices))) + if isinstance(positional.choices, dict): + public_cmds = get_public_subcommands(positional) + this_positional_choices = [] for choice in positional.choices: if isinstance(choice, Choice): @@ -222,7 +235,7 @@ def recurse(parser, prefix): elif isinstance(positional.choices, dict): # subparser, so append to list of subparsers & recurse log.debug("subcommand:%s", choice) - if positional.choices[choice].add_help: + if choice in public_cmds: discovered_subparsers.append(str(choice)) this_positional_choices.append(str(choice)) ( @@ -577,8 +590,9 @@ def format_positional(opt): root_arguments.append(format_positional(opt)) else: # subparser log.debug("choices:{}:{}".format(root_prefix, sorted(sub.choices))) + public_cmds = get_public_subcommands(sub) for cmd, subparser in sub.choices.items(): - if not subparser.add_help: + if cmd not in public_cmds: log.debug("skip:subcommand:%s", cmd) continue log.debug("subcommand:%s", cmd) diff --git a/tests/test_shtab.py b/tests/test_shtab.py index 5ce9959..be7f71b 100644 --- a/tests/test_shtab.py +++ b/tests/test_shtab.py @@ -177,7 +177,7 @@ def test_custom_complete(shell, caplog): def test_subparser_custom_complete(shell, caplog): parser = ArgumentParser(prog="test") subparsers = parser.add_subparsers() - sub = subparsers.add_parser("sub") + sub = subparsers.add_parser("sub", help="help message") sub.add_argument("posA").complete = {"bash": "_shtab_test_some_func"} preamble = {"bash": "_shtab_test_some_func() { compgen -W 'one two' -- $1 ;}"} with caplog.at_level(logging.INFO): @@ -194,6 +194,31 @@ def test_subparser_custom_complete(shell, caplog): assert not caplog.record_tuples +@fix_shell +def test_subparser_aliases(shell, caplog): + parser = ArgumentParser(prog="test") + subparsers = parser.add_subparsers() + sub = subparsers.add_parser("sub", aliases=["xsub", "ysub"], help="help message") + sub.add_argument("posA").complete = {"bash": "_shtab_test_some_func"} + preamble = {"bash": "_shtab_test_some_func() { compgen -W 'one two' -- $1 ;}"} + with caplog.at_level(logging.INFO): + completion = shtab.complete(parser, shell=shell, preamble=preamble) + print(completion) + + if shell == "bash": + shell = Bash(completion) + shell.compgen('-W "${_shtab_test_subparsers[*]}"', "s", "sub") + shell.compgen('-W "$_shtab_test_pos_0_choices"', "s", "sub") + shell.compgen('-W "${_shtab_test_subparsers[*]}"', "x", "xsub") + shell.compgen('-W "$_shtab_test_pos_0_choices"', "x", "xsub") + shell.compgen('-W "${_shtab_test_subparsers[*]}"', "y", "ysub") + shell.compgen('-W "$_shtab_test_pos_0_choices"', "y", "ysub") + shell.test('"$($_shtab_test_sub_pos_0_COMPGEN o)" = "one"') + shell.test('-z "$_shtab_test_COMPGEN"') + + assert not caplog.record_tuples + + @fix_shell def test_add_argument_to_optional(shell, caplog): parser = ArgumentParser(prog="test") @@ -213,7 +238,7 @@ def test_add_argument_to_optional(shell, caplog): def test_add_argument_to_positional(shell, caplog, capsys): parser = ArgumentParser(prog="test") subparsers = parser.add_subparsers() - sub = subparsers.add_parser("completion") + sub = subparsers.add_parser("completion", help="help message") shtab.add_argument_to(sub, "shell", parent=parser) from argparse import Namespace