Skip to content

Commit

Permalink
Determine whether to add a subparser to completion by its presence in…
Browse files Browse the repository at this point in the history
… _get_subactions()

This is the logic argparse uses for showing a subparser in the list of
options, so shtab now matches that behavior. This is particularly important
for subparsers which pass `add_help=False` in order to provide their *own*
help implementation.
  • Loading branch information
jimporter committed Aug 23, 2021
1 parent 76c5ead commit 546a840
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 4 deletions.
18 changes: 16 additions & 2 deletions shtab/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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))
(
Expand Down Expand Up @@ -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)
Expand Down
27 changes: 25 additions & 2 deletions tests/test_shtab.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -194,6 +194,29 @@ 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"', "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")
Expand All @@ -213,7 +236,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

Expand Down

0 comments on commit 546a840

Please sign in to comment.