Updated _set_parser_prog() so future calls to add_parser() will set the correct prog value

This makes dynamically adding subcommands after the CLI starts easier.

2 files changed, 13 insertions, 2 deletions

diff --git a/cmd2/decorators.py b/cmd2/decorators.py
index 7d097534..d2fdf9c7 100644
--- a/cmd2/decorators.py
+++ b/cmd2/decorators.py
@@ -87,7 +87,7 @@ def _set_parser_prog(parser: argparse.ArgumentParser, prog: str):
     is a command name and not sys.argv[0].
 
     :param parser: the parser being edited
-    :param prog: value for the current parsers prog attribute
+    :param prog: new value for the parser's prog attribute
     """
     # Set the prog value for this parser
     parser.prog = prog
@@ -95,6 +95,10 @@ def _set_parser_prog(parser: argparse.ArgumentParser, prog: str):
     # Set the prog value for the parser's subcommands
     for action in parser._actions:
         if isinstance(action, argparse._SubParsersAction):
+            # Set the _SubParsersAction's _prog_prefix value. That way if its add_parser() method is called later,
+            # the correct prog value will be set on the parser being added.
+            action._prog_prefix = parser.prog
+
             # The keys of action.choices are subcommand names as well as subcommand aliases. The aliases point to the
             # same parser as the actual subcommand. We want to avoid placing an alias into a parser's prog value.
             # Unfortunately there is nothing about an action.choices entry which tells us it's an alias. In most cases
diff --git a/tests/test_argparse.py b/tests/test_argparse.py
index f9ee0d07..db6dea21 100644
--- a/tests/test_argparse.py
+++ b/tests/test_argparse.py
@@ -333,7 +333,14 @@ def test_subcommand_help(subcommand_app):
     assert out[1] == ''
     assert out[2] == 'positional arguments:'
 
-
 def test_subcommand_invalid_help(subcommand_app):
     out, err = run_cmd(subcommand_app, 'help base baz')
     assert out[0].startswith('usage: base')
+
+def test_add_another_subcommand(subcommand_app):
+    """
+    This tests makes sure _set_parser_prog() sets _prog_prefix on every _SubParsersAction so that all future calls
+    to add_parser() write the correct prog value to the parser being added.
+    """
+    new_parser = subcommand_app.base_subparsers.add_parser('new_sub', help="stuff")
+    assert new_parser.prog == "base new_sub"