Adds support for injectable subcommands as part of CommandSet

load/unload.
Updated examples and documentation to include discussion of injectable
sub-commands.

5 files changed, 160 insertions, 8 deletions

diff --git a/cmd2/__init__.py b/cmd2/__init__.py
index 1fb01b16..19e620be 100644
--- a/cmd2/__init__.py
+++ b/cmd2/__init__.py
@@ -30,7 +30,7 @@ from .argparse_custom import DEFAULT_ARGUMENT_PARSER
 from .cmd2 import Cmd
 from .command_definition import CommandSet, with_default_category
 from .constants import COMMAND_NAME, DEFAULT_SHORTCUTS
-from .decorators import with_argument_list, with_argparser, with_argparser_and_unknown_args, with_category
+from .decorators import with_argument_list, with_argparser, with_argparser_and_unknown_args, with_category, as_subcommand_to
 from .exceptions import Cmd2ArgparseError, SkipPostcommandHooks
 from . import plugin
 from .parsing import Statement
diff --git a/cmd2/argparse_custom.py b/cmd2/argparse_custom.py
index 74bddfc7..689c1db7 100644
--- a/cmd2/argparse_custom.py
+++ b/cmd2/argparse_custom.py
@@ -724,6 +724,24 @@ class Cmd2HelpFormatter(argparse.RawTextHelpFormatter):
         return result
 
 
+class _UnloadableSubParsersAction(argparse._SubParsersAction):
+    """Extends the argparse internal SubParsers action to allow sub-parsers to be removed dynamically"""
+    def remove_parser(self, name):
+        """Removes a sub-parser from the sub-parsers group"""
+        for choice_action in self._choices_actions:
+            if choice_action.dest == name:
+                self._choices_actions.remove(choice_action)
+                break
+
+        subparser = self._name_parser_map[name]
+        to_remove = []
+        for name, parser in self._name_parser_map.items():
+            if parser is subparser:
+                to_remove.append(name)
+        for name in to_remove:
+            del self._name_parser_map[name]
+
+
 # noinspection PyCompatibility
 class Cmd2ArgumentParser(argparse.ArgumentParser):
     """Custom ArgumentParser class that improves error and help output"""
@@ -754,12 +772,22 @@ class Cmd2ArgumentParser(argparse.ArgumentParser):
             conflict_handler=conflict_handler,
             add_help=add_help,
             allow_abbrev=allow_abbrev)
+        self.register('action', 'unloadable_parsers', _UnloadableSubParsersAction)
+
+    def add_subparsers(self, unloadable=False, **kwargs):
+        """
+        Custom override. Sets a default title if one was not given.
 
-    def add_subparsers(self, **kwargs):
-        """Custom override. Sets a default title if one was not given."""
+        :param unloadable: Flag whether this sub-parsers group will support unloading parsers
+        :param kwargs: additional keyword arguments
+        :return: argparse Subparser Action
+        """
         if 'title' not in kwargs:
             kwargs['title'] = 'subcommands'
 
+        if unloadable:
+            kwargs['action'] = 'unloadable_parsers'
+
         return super().add_subparsers(**kwargs)
 
     def error(self, message: str) -> None:
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index ca60a461..2a1bcbf2 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -43,7 +43,7 @@ from contextlib import redirect_stdout
 from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Type, Union
 
 from . import ansi, constants, plugin, utils
-from .argparse_custom import DEFAULT_ARGUMENT_PARSER, CompletionItem
+from .argparse_custom import DEFAULT_ARGUMENT_PARSER, CompletionItem, _UnloadableSubParsersAction
 from .clipboard import can_clip, get_paste_buffer, write_to_paste_buffer
 from .command_definition import CommandSet, _partial_passthru
 from .constants import COMMAND_FUNC_PREFIX, COMPLETER_FUNC_PREFIX, HELP_FUNC_PREFIX
@@ -260,6 +260,8 @@ class Cmd(cmd.Cmd):
             if not valid:
                 raise ValueError("Invalid command name {!r}: {}".format(cur_cmd, errmsg))
 
+        self._register_subcommands(self)
+
         # Stores results from the last command run to enable usage of results in a Python script or interactive console
         # Built-in commands don't make use of this.  It is purely there for user-defined commands and convenience.
         self.last_result = None
@@ -429,12 +431,12 @@ class Cmd(cmd.Cmd):
 
         installed_attributes = []
         try:
-            for method in methods:
-                command = method[0][len(COMMAND_FUNC_PREFIX):]
-                command_wrapper = _partial_passthru(method[1], self)
+            for method_name, method in methods:
+                command = method_name[len(COMMAND_FUNC_PREFIX):]
+                command_wrapper = _partial_passthru(method, self)
 
                 self.__install_command_function(command, command_wrapper, type(cmdset).__name__)
-                installed_attributes.append(method[0])
+                installed_attributes.append(method_name)
 
                 completer_func_name = COMPLETER_FUNC_PREFIX + command
                 cmd_completer = getattr(cmdset, completer_func_name, None)
@@ -451,6 +453,8 @@ class Cmd(cmd.Cmd):
                     installed_attributes.append(help_func_name)
 
             self._installed_command_sets.append(cmdset)
+
+            self._register_subcommands(cmdset)
         except Exception:
             for attrib in installed_attributes:
                 delattr(self, attrib)
@@ -500,6 +504,9 @@ class Cmd(cmd.Cmd):
         :param cmdset: CommandSet to uninstall
         """
         if cmdset in self._installed_command_sets:
+
+            self._unregister_subcommands(cmdset)
+
             methods = inspect.getmembers(
                 cmdset,
                 predicate=lambda meth: inspect.ismethod(meth) and meth.__name__.startswith(COMMAND_FUNC_PREFIX))
@@ -517,6 +524,88 @@ class Cmd(cmd.Cmd):
             cmdset.on_unregister(self)
             self._installed_command_sets.remove(cmdset)
 
+    def _register_subcommands(self, cmdset: Union[CommandSet, 'Cmd']) -> None:
+        """
+        Register sub-commands with their base command
+
+        :param cmdset: CommandSet containing sub-commands
+        """
+        if not (cmdset is self or cmdset in self._installed_command_sets):
+            raise ValueError('Adding sub-commands from an unregistered CommandSet')
+
+        # find all methods that start with the sub-command prefix
+        methods = inspect.getmembers(
+            cmdset,
+            predicate=lambda meth: (inspect.ismethod(meth) or isinstance(meth, Callable))
+            and hasattr(meth, constants.SUBCMD_ATTR_NAME)
+            and hasattr(meth, constants.SUBCMD_ATTR_COMMAND)
+            and hasattr(meth, constants.CMD_ATTR_ARGPARSER)
+        )
+
+        # iterate through all matching methods
+        for method_name, method in methods:
+            subcommand_name = getattr(method, constants.SUBCMD_ATTR_NAME)
+            command_name = getattr(method, constants.SUBCMD_ATTR_COMMAND)
+            subcmd_parser = getattr(method, constants.CMD_ATTR_ARGPARSER)
+
+            # Search for the base command function and verify it has an argparser defined
+            command_func = self.cmd_func(command_name)
+            if command_func is None or not hasattr(command_func, constants.CMD_ATTR_ARGPARSER):
+                raise TypeError('Could not find command: ' + command_name + ' needed by sub-command ' + str(method))
+            command_parser = getattr(command_func, constants.CMD_ATTR_ARGPARSER)
+            if command_parser is None:
+                raise TypeError('Could not find argparser for command: ' + command_name + ' needed by sub-command ' + str(method))
+
+            if hasattr(method, '__doc__') and method.__doc__ is not None:
+                help_text = method.__doc__.splitlines()[0]
+            else:
+                help_text = subcommand_name
+
+            if isinstance(cmdset, CommandSet):
+                command_handler = _partial_passthru(method, self)
+            else:
+                command_handler = method
+            subcmd_parser.set_defaults(handler=command_handler)
+
+            for action in command_parser._actions:
+                if isinstance(action, _UnloadableSubParsersAction):
+                    action.add_parser(subcommand_name, parents=[subcmd_parser], help=help_text)
+
+    def _unregister_subcommands(self, cmdset: Union[CommandSet, 'Cmd']) -> None:
+        """
+        Unregister sub-commands from their base command
+
+        :param cmdset: CommandSet containing sub-commands
+        """
+        if not (cmdset is self or cmdset in self._installed_command_sets):
+            raise ValueError('Removing sub-commands from an unregistered CommandSet')
+
+        # find all methods that start with the sub-command prefix
+        methods = inspect.getmembers(
+            cmdset,
+            predicate=lambda meth: (inspect.ismethod(meth) or isinstance(meth, Callable))
+            and hasattr(meth, constants.SUBCMD_ATTR_NAME)
+            and hasattr(meth, constants.SUBCMD_ATTR_COMMAND)
+            and hasattr(meth, constants.CMD_ATTR_ARGPARSER)
+        )
+
+        # iterate through all matching methods
+        for method_name, method in methods:
+            subcommand_name = getattr(method, constants.SUBCMD_ATTR_NAME)
+            command_name = getattr(method, constants.SUBCMD_ATTR_COMMAND)
+
+            # Search for the base command function and verify it has an argparser defined
+            command_func = self.cmd_func(command_name)
+            if command_func is None or not hasattr(command_func, constants.CMD_ATTR_ARGPARSER):
+                raise TypeError('Could not find command: ' + command_name + ' needed by sub-command ' + str(method))
+            command_parser = getattr(command_func, constants.CMD_ATTR_ARGPARSER)
+            if command_parser is None:
+                raise TypeError('Could not find argparser for command: ' + command_name + ' needed by sub-command ' + str(method))
+
+            for action in command_parser._actions:
+                if isinstance(action, _UnloadableSubParsersAction):
+                    action.remove_parser(subcommand_name)
+
     def add_settable(self, settable: Settable) -> None:
         """
         Convenience method to add a settable parameter to ``self.settables``
diff --git a/cmd2/constants.py b/cmd2/constants.py
index 81d1a29b..0135e328 100644
--- a/cmd2/constants.py
+++ b/cmd2/constants.py
@@ -49,3 +49,7 @@ CMD_ATTR_ARGPARSER = 'argparser'
 
 # Whether or not tokens are unquoted before sending to argparse
 CMD_ATTR_PRESERVE_QUOTES = 'preserve_quotes'
+
+# sub-command attributes for the base command name and the sub-command name
+SUBCMD_ATTR_COMMAND = 'parent_command'
+SUBCMD_ATTR_NAME = 'subcommand_name'
diff --git a/cmd2/decorators.py b/cmd2/decorators.py
index 8c3739f1..6e3b7acf 100644
--- a/cmd2/decorators.py
+++ b/cmd2/decorators.py
@@ -335,3 +335,34 @@ def with_argparser(parser: argparse.ArgumentParser, *,
 
     # noinspection PyTypeChecker
     return arg_decorator
+
+
+def as_subcommand_to(command: str,
+                     subcommand: str,
+                     parser: argparse.ArgumentParser) -> Callable[[argparse.Namespace], Optional[bool]]:
+    """
+    Tag this method as a sub-command to an existing argparse decorated command.
+
+    :param command: Command Name
+    :param subcommand: Sub-command name
+    :param parser: argparse Parser to for this sub-command
+    :return: Wrapper function that can receive an argparse.Namespace
+    """
+    def arg_decorator(func: Callable):
+        _set_parser_prog(parser, subcommand)
+
+        # If the description has not been set, then use the method docstring if one exists
+        if parser.description is None and func.__doc__:
+            parser.description = func.__doc__
+
+        parser.set_defaults(func=func)
+
+        # # Set some custom attributes for this command
+        setattr(func, constants.SUBCMD_ATTR_COMMAND, command)
+        setattr(func, constants.CMD_ATTR_ARGPARSER, parser)
+        setattr(func, constants.SUBCMD_ATTR_NAME, subcommand)
+
+        return func
+
+    # noinspection PyTypeChecker
+    return arg_decorator