Repository: iterative/shtab

Issue:
subcommand convenience wrapper doesn't work
https://github.com/iterative/shtab/blob/a01e2b592102496a3a7b6861b448d0d6e78407f4/examples/customcomplete.py#L43 provides completion only for the `subparser` rather than the full (parent) `parser`

Discussion:

Patch:
diff --git a/examples/customcomplete.py b/examples/customcomplete.py
index f539187..aa24d4e 100755
--- a/examples/customcomplete.py
+++ b/examples/customcomplete.py
@@ -40,7 +40,7 @@ def get_main_parser():
     subparsers.dest = "subcommand"
 
     parser = subparsers.add_parser("completion")
-    shtab.add_argument_to(parser, "shell")  # magic!
+    shtab.add_argument_to(parser, "shell", parent=main_parser)  # magic!
 
     parser = subparsers.add_parser("process")
     # `*.txt` file tab completion
diff --git a/shtab/__init__.py b/shtab/__init__.py
index 3592d89..25e1f09 100644
--- a/shtab/__init__.py
+++ b/shtab/__init__.py
@@ -602,22 +602,27 @@ def complete(
     )
 
 
-class PrintCompletionAction(Action):
-    def __call__(self, parser, namespace, values, option_string=None):
-        print(complete(parser, values))
-        parser.exit(0)
+def completion_action(parent=None):
+    class PrintCompletionAction(Action):
+        def __call__(self, parser, namespace, values, option_string=None):
+            print(complete(parent or parser, values))
+            parser.exit(0)
+
+    return PrintCompletionAction
 
 
 def add_argument_to(
     parser,
     option_string="--print-completion",
     help="print shell completion script",
+    parent=None,
 ):
     """
     parser  : argparse.ArgumentParser
     option_string  : str or list[str], iff positional (no `-` prefix) then
       `parser` is assumed to actually be a subparser (subcommand mode)
     help  : str
+    parent  : argparse.ArgumentParser, required in subcommand mode
     """
     if isinstance(
         option_string, str if sys.version_info[0] > 2 else basestring  # NOQA
@@ -627,9 +632,10 @@ def add_argument_to(
         choices=SUPPORTED_SHELLS,
         default=None,
         help=help,
-        action=PrintCompletionAction,
+        action=completion_action(parent),
     )
     if option_string[0][0] != "-":  # subparser mode
         kwargs.update(default=SUPPORTED_SHELLS[0], nargs="?")
+        assert parent is not None, "subcommand mode: parent required"
     parser.add_argument(*option_string, **kwargs)
     return parser


Test patch:
diff --git a/tests/test_shtab.py b/tests/test_shtab.py
index 54625ab..2819848 100644
--- a/tests/test_shtab.py
+++ b/tests/test_shtab.py
@@ -214,17 +214,27 @@ def test_add_argument_to_optional(shell, caplog):
 
 
 @fix_shell
-def test_add_argument_to_positional(shell, caplog):
+def test_add_argument_to_positional(shell, caplog, capsys):
     parser = ArgumentParser(prog="test")
     subparsers = parser.add_subparsers()
     sub = subparsers.add_parser("completion")
-    shtab.add_argument_to(sub, "shell")
+    shtab.add_argument_to(sub, "shell", parent=parser)
+    from argparse import Namespace
+
     with caplog.at_level(logging.INFO):
-        completion = shtab.complete(parser, shell=shell)
+        completion_manual = shtab.complete(parser, shell=shell)
+        with pytest.raises(SystemExit) as exc:
+            sub._actions[-1](sub, Namespace(), shell)
+            assert exc.type == SystemExit
+            assert exc.vaue.code == 0
+    completion, err = capsys.readouterr()
     print(completion)
+    assert completion_manual == completion.rstrip()
+    assert not err
 
     if shell == "bash":
         shell = Bash(completion)
+        shell.compgen('-W "$_shtab_test_commands_"', "c", "completion")
         shell.compgen('-W "$_shtab_test_completion"', "ba", "bash")
         shell.compgen('-W "$_shtab_test_completion"', "z", "zsh")