Refactored custom ArgparseCompleter functionality so they will now be set using methods on ArgumentParser objects.

This fixes issue where subcommands did not use the correct custom ArgparseCompleter type.

6 files changed, 121 insertions, 54 deletions

diff --git a/cmd2/__init__.py b/cmd2/__init__.py
index a23b7a36..16158900 100644
--- a/cmd2/__init__.py
+++ b/cmd2/__init__.py
@@ -25,7 +25,7 @@ from .argparse_custom import (
     Cmd2AttributeWrapper,
     CompletionItem,
     register_argparse_argument_parameter,
-    set_default_argument_parser,
+    set_default_argument_parser_type,
 )
 
 # Check if user has defined a module that sets a custom value for argparse_custom.DEFAULT_ARGUMENT_PARSER.
@@ -38,7 +38,7 @@ if cmd2_parser_module is not None:
 
     importlib.import_module(cmd2_parser_module)
 
-from .argparse_completer import set_default_command_completer_type
+from .argparse_completer import set_default_ap_completer_type
 
 from .cmd2 import Cmd
 from .command_definition import CommandSet, with_default_category
@@ -63,8 +63,8 @@ __all__: List[str] = [
     'Cmd2AttributeWrapper',
     'CompletionItem',
     'register_argparse_argument_parameter',
-    'set_default_argument_parser',
-    'set_default_command_completer_type',
+    'set_default_argument_parser_type',
+    'set_default_ap_completer_type',
     # Cmd2
     'Cmd',
     'CommandResult',
diff --git a/cmd2/argparse_completer.py b/cmd2/argparse_completer.py
index ebc49a8c..f31584e7 100644
--- a/cmd2/argparse_completer.py
+++ b/cmd2/argparse_completer.py
@@ -407,9 +407,15 @@ class ArgparseCompleter:
                                 if action.dest != argparse.SUPPRESS:
                                     parent_tokens[action.dest] = [token]
 
-                                completer = ArgparseCompleter(
-                                    self._subcommand_action.choices[token], self._cmd2_app, parent_tokens=parent_tokens
-                                )
+                                parser: argparse.ArgumentParser = self._subcommand_action.choices[token]
+                                completer_type: Optional[
+                                    Type[ArgparseCompleter]
+                                ] = parser.get_ap_completer_type()  # type: ignore[attr-defined]
+                                if completer_type is None:
+                                    completer_type = DEFAULT_AP_COMPLETER
+
+                                completer = completer_type(parser, self._cmd2_app, parent_tokens=parent_tokens)
+
                                 return completer.complete(
                                     text, line, begidx, endidx, tokens[token_index + 1 :], cmd_set=cmd_set
                                 )
@@ -609,7 +615,14 @@ class ArgparseCompleter:
         if self._subcommand_action is not None:
             for token_index, token in enumerate(tokens):
                 if token in self._subcommand_action.choices:
-                    completer = ArgparseCompleter(self._subcommand_action.choices[token], self._cmd2_app)
+                    parser: argparse.ArgumentParser = self._subcommand_action.choices[token]
+                    completer_type: Optional[
+                        Type[ArgparseCompleter]
+                    ] = parser.get_ap_completer_type()  # type: ignore[attr-defined]
+                    if completer_type is None:
+                        completer_type = DEFAULT_AP_COMPLETER
+
+                    completer = completer_type(parser, self._cmd2_app)
                     return completer.complete_subcommand_help(text, line, begidx, endidx, tokens[token_index + 1 :])
                 elif token_index == len(tokens) - 1:
                     # Since this is the last token, we will attempt to complete it
@@ -629,7 +642,14 @@ class ArgparseCompleter:
         if self._subcommand_action is not None:
             for token_index, token in enumerate(tokens):
                 if token in self._subcommand_action.choices:
-                    completer = ArgparseCompleter(self._subcommand_action.choices[token], self._cmd2_app)
+                    parser: argparse.ArgumentParser = self._subcommand_action.choices[token]
+                    completer_type: Optional[
+                        Type[ArgparseCompleter]
+                    ] = parser.get_ap_completer_type()  # type: ignore[attr-defined]
+                    if completer_type is None:
+                        completer_type = DEFAULT_AP_COMPLETER
+
+                    completer = completer_type(parser, self._cmd2_app)
                     return completer.format_help(tokens[token_index + 1 :])
                 else:
                     break
@@ -740,14 +760,15 @@ class ArgparseCompleter:
         return self._format_completions(arg_state, results)
 
 
-DEFAULT_COMMAND_COMPLETER: Type[ArgparseCompleter] = ArgparseCompleter
+# The default ArgparseCompleter class for a cmd2 app
+DEFAULT_AP_COMPLETER: Type[ArgparseCompleter] = ArgparseCompleter
 
 
-def set_default_command_completer_type(completer_type: Type[ArgparseCompleter]) -> None:
+def set_default_ap_completer_type(completer_type: Type[ArgparseCompleter]) -> None:
     """
-    Set the default command completer type. It must be a sub-class of the ArgparseCompleter.
+    Set the default ArgparseCompleter class for a cmd2 app.
 
     :param completer_type: Type that is a subclass of ArgparseCompleter.
     """
-    global DEFAULT_COMMAND_COMPLETER
-    DEFAULT_COMMAND_COMPLETER = completer_type
+    global DEFAULT_AP_COMPLETER
+    DEFAULT_AP_COMPLETER = completer_type
diff --git a/cmd2/argparse_custom.py b/cmd2/argparse_custom.py
index 44e7a90b..dd4db570 100644
--- a/cmd2/argparse_custom.py
+++ b/cmd2/argparse_custom.py
@@ -207,6 +207,13 @@ for cases in which you need to manually access the cmd2-specific attributes.
 - ``argparse.Action.set_suppress_tab_hint()`` - See
   :func:`_action_set_suppress_tab_hint` for more details.
 
+cmd2 has patched ``argparse.ArgumentParser`` to include the following accessor methods
+
+- ``argparse.ArgumentParser.get_ap_completer_type()`` - See
+  :func:`_ArgumentParser_get_ap_completer_type` for more details.
+- ``argparse.Action.set_ap_completer_type()`` - See
+  :func:`_ArgumentParser_set_ap_completer_type` for more details.
+
 **Subcommand removal**
 
 cmd2 has patched ``argparse._SubParsersAction`` to include a ``remove_parser()``
@@ -232,6 +239,7 @@ from gettext import (
 )
 from typing import (
     IO,
+    TYPE_CHECKING,
     Any,
     Callable,
     Dict,
@@ -264,6 +272,12 @@ except ImportError:
     )
 
 
+if TYPE_CHECKING:  # pragma: no cover
+    from .argparse_completer import (
+        ArgparseCompleter,
+    )
+
+
 def generate_range_error(range_min: int, range_max: Union[int, float]) -> str:
     """Generate an error message when the the number of arguments provided is not within the expected range"""
     err_str = "expected "
@@ -659,6 +673,7 @@ def register_argparse_argument_parameter(param_name: str, param_type: Optional[T
     and ``set_{param_name}(value)``.
 
     :param param_name: Name of the parameter to add.
+    :param param_type: Type of the parameter to add.
     """
     attr_name = f'{_CUSTOM_ATTRIB_PFX}{param_name}'
     if param_name in CUSTOM_ACTION_ATTRIBS or hasattr(argparse.Action, attr_name):
@@ -715,6 +730,7 @@ def register_argparse_argument_parameter(param_name: str, param_type: Optional[T
 orig_actions_container_add_argument = argparse._ActionsContainer.add_argument
 
 
+# noinspection PyProtectedMember
 def _add_argument_wrapper(
     self: argparse._ActionsContainer,
     *args: Any,
@@ -916,10 +932,54 @@ setattr(argparse.ArgumentParser, '_match_argument', _match_argument_wrapper)
 
 
 ############################################################################################################
-# Patch argparse._SubParsersAction to add remove_parser function
+# Patch argparse.ArgumentParser with accessors for ap_completer_type attribute
 ############################################################################################################
 
+# An ArgumentParser attribute which specifies a subclass of ArgparseCompleter for custom tab completion behavior on a
+# given parser. If this is None or not present, then cmd2 will use argparse_completer.DEFAULT_AP_COMPLETER when tab
+# completing a parser's arguments
+ATTR_AP_COMPLETER_TYPE = 'ap_completer_type'
+
+
 # noinspection PyPep8Naming
+def _ArgumentParser_get_ap_completer_type(self: argparse.ArgumentParser) -> Optional[Type['ArgparseCompleter']]:
+    """
+    Get the ap_completer_type attribute of an argparse ArgumentParser.
+
+    This function is added by cmd2 as a method called ``get_ap_completer_type()`` to ``argparse.ArgumentParser`` class.
+
+    To call: ``parser.get_ap_completer_type()``
+
+    :param self: ArgumentParser being queried
+    :return: An ArgparseCompleter-based class or None if attribute does not exist
+    """
+    return cast(Optional[Type['ArgparseCompleter']], getattr(self, ATTR_AP_COMPLETER_TYPE, None))
+
+
+setattr(argparse.ArgumentParser, 'get_ap_completer_type', _ArgumentParser_get_ap_completer_type)
+
+
+# noinspection PyPep8Naming
+def _ArgumentParser_set_ap_completer_type(self: argparse.ArgumentParser, ap_completer_type: Type['ArgparseCompleter']) -> None:
+    """
+    Set the ap_completer_type attribute of an argparse ArgumentParser.
+
+    This function is added by cmd2 as a method called ``set_ap_completer_type()`` to ``argparse.ArgumentParser`` class.
+
+    :param self: ArgumentParser being edited
+    :param ap_completer_type: the custom ArgparseCompleter-based class to use when tab completing arguments for this parser
+    """
+    setattr(self, ATTR_AP_COMPLETER_TYPE, ap_completer_type)
+
+
+setattr(argparse.ArgumentParser, 'set_ap_completer_type', _ArgumentParser_set_ap_completer_type)
+
+
+############################################################################################################
+# Patch argparse._SubParsersAction to add remove_parser function
+############################################################################################################
+
+# noinspection PyPep8Naming,PyProtectedMember
 def _SubParsersAction_remove_parser(self: argparse._SubParsersAction, name: str) -> None:
     """
     Removes a sub-parser from a sub-parsers group. Used to remove subcommands from a parser.
@@ -964,6 +1024,7 @@ setattr(argparse._SubParsersAction, 'remove_parser', _SubParsersAction_remove_pa
 class Cmd2HelpFormatter(argparse.RawTextHelpFormatter):
     """Custom help formatter to configure ordering of help text"""
 
+    # noinspection PyProtectedMember
     def _format_usage(
         self,
         usage: Optional[str],
@@ -1207,6 +1268,7 @@ class Cmd2ArgumentParser(argparse.ArgumentParser):
             allow_abbrev=allow_abbrev,
         )
 
+    # noinspection PyProtectedMember
     def add_subparsers(self, **kwargs: Any) -> argparse._SubParsersAction:
         """
         Custom override. Sets a default title if one was not given.
@@ -1321,10 +1383,10 @@ class Cmd2AttributeWrapper:
 DEFAULT_ARGUMENT_PARSER: Type[argparse.ArgumentParser] = Cmd2ArgumentParser
 
 
-def set_default_argument_parser(parser: Type[argparse.ArgumentParser]) -> None:
+def set_default_argument_parser_type(parser_type: Type[argparse.ArgumentParser]) -> None:
     """
     Set the default ArgumentParser class for a cmd2 app. This must be called prior to loading cmd2.py if
     you want to override the parser for cmd2's built-in commands. See examples/override_parser.py.
     """
     global DEFAULT_ARGUMENT_PARSER
-    DEFAULT_ARGUMENT_PARSER = parser
+    DEFAULT_ARGUMENT_PARSER = parser_type
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 1516f709..128cfa76 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -860,6 +860,9 @@ class Cmd(cmd.Cmd):
                     defaults = {constants.NS_ATTR_SUBCMD_HANDLER: method}
                     attached_parser.set_defaults(**defaults)
 
+                    # Copy value for custom ArgparseCompleter type, which will be None if not present on subcmd_parser
+                    attached_parser.set_ap_completer_type(subcmd_parser.get_ap_completer_type())  # type: ignore[attr-defined]
+
                     # Set what instance the handler is bound to
                     setattr(attached_parser, constants.PARSER_ATTR_COMMANDSET, cmdset)
                     break
@@ -1850,10 +1853,6 @@ class Cmd(cmd.Cmd):
         :param endidx: the ending index of the prefix text
         :param custom_settings: optional prepopulated completion settings
         """
-        from .argparse_completer import (
-            ArgparseCompleter,
-        )
-
         # If custom_settings is None, then we are completing a command's argument.
         # Parse the command line to get the command token.
         command = ''
@@ -1903,18 +1902,18 @@ class Cmd(cmd.Cmd):
                 else:
                     # There's no completer function, next see if the command uses argparse
                     func = self.cmd_func(command)
-                    argparser = getattr(func, constants.CMD_ATTR_ARGPARSER, None)
-                    completer_type = getattr(func, constants.CMD_ATTR_COMPLETER, argparse_completer.DEFAULT_COMMAND_COMPLETER)
-                    if completer_type is None:
-                        completer_type = argparse_completer.DEFAULT_COMMAND_COMPLETER
+                    argparser: Optional[argparse.ArgumentParser] = getattr(func, constants.CMD_ATTR_ARGPARSER, None)
 
                     if func is not None and argparser is not None:
-                        cmd_set = self._cmd_to_command_sets[command] if command in self._cmd_to_command_sets else None
-                        if completer_type is not None:
-                            completer = completer_type(argparser, self)
-                        else:
-                            completer = ArgparseCompleter(argparser, self)
+                        # Get arguments for complete()
                         preserve_quotes = getattr(func, constants.CMD_ATTR_PRESERVE_QUOTES)
+                        cmd_set = self._cmd_to_command_sets[command] if command in self._cmd_to_command_sets else None
+
+                        # Create the argparse completer
+                        completer_type = argparser.get_ap_completer_type()  # type: ignore[attr-defined]
+                        if completer_type is None:
+                            completer_type = argparse_completer.DEFAULT_AP_COMPLETER
+                        completer = completer_type(argparser, self)
 
                         completer_func = functools.partial(
                             completer.complete, tokens=raw_tokens[1:] if preserve_quotes else tokens[1:], cmd_set=cmd_set
@@ -1932,7 +1931,12 @@ class Cmd(cmd.Cmd):
 
         # Otherwise we are completing the command token or performing custom completion
         else:
-            completer = ArgparseCompleter(custom_settings.parser, self)
+            # Create the argparse completer
+            completer_type = custom_settings.parser.get_ap_completer_type()  # type: ignore[attr-defined]
+            if completer_type is None:
+                completer_type = argparse_completer.DEFAULT_AP_COMPLETER
+            completer = completer_type(custom_settings.parser, self)
+
             completer_func = functools.partial(
                 completer.complete, tokens=raw_tokens if custom_settings.preserve_quotes else tokens, cmd_set=None
             )
@@ -3542,11 +3546,7 @@ class Cmd(cmd.Cmd):
         if func is None or argparser is None:
             return []
 
-        from .argparse_completer import (
-            ArgparseCompleter,
-        )
-
-        completer = ArgparseCompleter(argparser, self)
+        completer = argparse_completer.DEFAULT_AP_COMPLETER(argparser, self)
         return completer.complete_subcommand_help(text, line, begidx, endidx, arg_tokens['subcommands'])
 
     help_parser = argparse_custom.DEFAULT_ARGUMENT_PARSER(
@@ -3582,11 +3582,7 @@ class Cmd(cmd.Cmd):
 
             # If the command function uses argparse, then use argparse's help
             if func is not None and argparser is not None:
-                from .argparse_completer import (
-                    ArgparseCompleter,
-                )
-
-                completer = ArgparseCompleter(argparser, self)
+                completer = argparse_completer.DEFAULT_AP_COMPLETER(argparser, self)
 
                 # Set end to blank so the help output matches how it looks when "command -h" is used
                 self.poutput(completer.format_help(args.subcommands), end='')
@@ -3918,11 +3914,7 @@ class Cmd(cmd.Cmd):
             completer=settable.completer,
         )
 
-        from .argparse_completer import (
-            ArgparseCompleter,
-        )
-
-        completer = ArgparseCompleter(settable_parser, self)
+        completer = argparse_completer.DEFAULT_AP_COMPLETER(settable_parser, self)
 
         # Use raw_tokens since quotes have been preserved
         _, raw_tokens = self.tokens_for_completion(line, begidx, endidx)
diff --git a/cmd2/constants.py b/cmd2/constants.py
index 7656ae58..9f29be86 100644
--- a/cmd2/constants.py
+++ b/cmd2/constants.py
@@ -43,7 +43,6 @@ CLASS_ATTR_DEFAULT_HELP_CATEGORY = 'cmd2_default_help_category'
 
 # The argparse parser for the command
 CMD_ATTR_ARGPARSER = 'argparser'
-CMD_ATTR_COMPLETER = 'command_completer'
 
 # Whether or not tokens are unquoted before sending to argparse
 CMD_ATTR_PRESERVE_QUOTES = 'preserve_quotes'
diff --git a/cmd2/decorators.py b/cmd2/decorators.py
index 644a8add..1ff0bdbe 100644
--- a/cmd2/decorators.py
+++ b/cmd2/decorators.py
@@ -10,16 +10,12 @@ from typing import (
     Optional,
     Sequence,
     Tuple,
-    Type,
     Union,
 )
 
 from . import (
     constants,
 )
-from .argparse_completer import (
-    ArgparseCompleter,
-)
 from .argparse_custom import (
     Cmd2AttributeWrapper,
 )
@@ -275,7 +271,6 @@ def with_argparser(
     ns_provider: Optional[Callable[..., argparse.Namespace]] = None,
     preserve_quotes: bool = False,
     with_unknown_args: bool = False,
-    completer: Optional[Type[ArgparseCompleter]] = None,
 ) -> Callable[[ArgparseCommandFunc], RawCommandFuncOptionalBoolReturn]:
     """A decorator to alter a cmd2 method to populate its ``args`` argument by parsing arguments
     with the given instance of argparse.ArgumentParser.
@@ -286,7 +281,6 @@ def with_argparser(
                         state data that affects parsing.
     :param preserve_quotes: if ``True``, then arguments passed to argparse maintain their quotes
     :param with_unknown_args: if true, then capture unknown args
-    :param completer: CommandCompleter type. Defaults to ArgparseCompleter if unspecified.
     :return: function that gets passed argparse-parsed args in a ``Namespace``
              A :class:`cmd2.argparse_custom.Cmd2AttributeWrapper` called ``cmd2_statement`` is included
              in the ``Namespace`` to provide access to the :class:`cmd2.Statement` object that was created when
@@ -397,7 +391,6 @@ def with_argparser(
 
         # Set some custom attributes for this command
         setattr(cmd_wrapper, constants.CMD_ATTR_ARGPARSER, parser)
-        setattr(cmd_wrapper, constants.CMD_ATTR_COMPLETER, completer)
         setattr(cmd_wrapper, constants.CMD_ATTR_PRESERVE_QUOTES, preserve_quotes)
 
         return cmd_wrapper