When passing a ns_provider to an argparse command, will now attempt to resolve the correct CommandSet instance for self. If not, it'll fall back and pass in the cmd2 app

4 files changed, 65 insertions, 44 deletions

diff --git a/cmd2/argparse_completer.py b/cmd2/argparse_completer.py
index 77fa41b8..582f57f6 100644
--- a/cmd2/argparse_completer.py
+++ b/cmd2/argparse_completer.py
@@ -569,45 +569,15 @@ class ArgparseCompleter:
         kwargs = {}
         if isinstance(arg_choices, ChoicesCallable):
             if arg_choices.is_method:
-                # figure out what class the completer was defined in
-                completer_class = get_defining_class(arg_choices.to_call)
-
-                # Was there a defining class identified? If so, is it a sub-class of CommandSet?
-                if completer_class is not None and issubclass(completer_class, CommandSet):
-                    # Since the completer function is provided as an unbound function, we need to locate the instance
-                    # of the CommandSet to pass in as `self` to emulate a bound method call.
-                    # We're searching for candidates that match the completer function's parent type in this order:
-                    #   1. Does the CommandSet registered with the command's argparser match as a subclass?
-                    #   2. Do any of the registered CommandSets in the Cmd2 application exactly match the type?
-                    #   3. Is there a registered CommandSet that is is the only matching subclass?
-
-                    # Now get the CommandSet associated with the current command/subcommand argparser
-                    parser_cmd_set = getattr(self._parser, constants.PARSER_ATTR_COMMANDSET, cmd_set)
-                    if isinstance(parser_cmd_set, completer_class):
-                        # Case 1: Parser's CommandSet is a sub-class of the completer function's CommandSet
-                        cmd_set = parser_cmd_set
-                    else:
-                        # Search all registered CommandSets
-                        cmd_set = None
-                        candidate_sets = []  # type: List[CommandSet]
-                        for installed_cmd_set in self._cmd2_app._installed_command_sets:
-                            if type(installed_cmd_set) == completer_class:
-                                # Case 2: CommandSet is an exact type match for the completer's CommandSet
-                                cmd_set = installed_cmd_set
-                                break
-
-                            # Add candidate for Case 3:
-                            if isinstance(installed_cmd_set, completer_class):
-                                candidate_sets.append(installed_cmd_set)
-                        if cmd_set is None and len(candidate_sets) == 1:
-                            # Case 3: There exists exactly 1 CommandSet that is a subclass of the completer's CommandSet
-                            cmd_set = candidate_sets[0]
-                    if cmd_set is None:
-                        # No cases matched, raise an error
-                        raise CompletionError('Could not find CommandSet instance matching defining type for completer')
-                    args.append(cmd_set)
-                else:
-                    args.append(self._cmd2_app)
+                # The completer may or may not be defined in the same class as the command. Since completer
+                # functions are registered with the command argparser before anything is instantiated, we
+                # need to find an instance at runtime that matches the types during declaration
+                cmd_set = self._cmd2_app._resolve_func_self(arg_choices.to_call, cmd_set)
+                if cmd_set is None:
+                    # No cases matched, raise an error
+                    raise CompletionError('Could not find CommandSet instance matching defining type for completer')
+
+                args.append(cmd_set)
 
             # Check if arg_choices.to_call expects arg_tokens
             to_call_params = inspect.signature(arg_choices.to_call).parameters
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 30dcb6e8..610ce4a3 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -59,7 +59,7 @@ from .exceptions import (
 from .history import History, HistoryItem
 from .parsing import Macro, MacroArg, Statement, StatementParser, shlex_split
 from .rl_utils import RlType, rl_get_point, rl_make_safe_prompt, rl_set_prompt, rl_type, rl_warning, vt100_support
-from .utils import CompletionError, Settable
+from .utils import CompletionError, get_defining_class, Settable
 
 # Set up readline
 if rl_type == RlType.NONE:  # pragma: no cover
@@ -4656,3 +4656,51 @@ class Cmd(cmd.Cmd):
         """Register a hook to be called after a command is completed, whether it completes successfully or not."""
         self._validate_cmdfinalization_callable(func)
         self._cmdfinalization_hooks.append(func)
+
+    def _resolve_func_self(self,
+                           cmd_support_func: Callable,
+                           cmd_self: Union[CommandSet, 'Cmd']) -> object:
+        """
+        Attempt to resolve a candidate instance to pass as 'self' for an unbound class method that was
+        used when defining command's argparse object. Since we restrict registration to only a single CommandSet
+        instance of each type, using type is a reasonably safe way to resolve the correct object instance
+
+        :param cmd_support_func: command support function. This could be a completer or namespace provider
+        :param cmd_self: The `self` associated with the command or sub-command
+        :return:
+        """
+        # figure out what class the command support function was defined in
+        func_class = get_defining_class(cmd_support_func)
+
+        # Was there a defining class identified? If so, is it a sub-class of CommandSet?
+        if func_class is not None and issubclass(func_class, CommandSet):
+            # Since the support function is provided as an unbound function, we need to locate the instance
+            # of the CommandSet to pass in as `self` to emulate a bound method call.
+            # We're searching for candidates that match the support function's defining class type in this order:
+            #   1. Is the command's CommandSet a sub-class of the support function's class?
+            #   2. Do any of the registered CommandSets in the Cmd2 application exactly match the type?
+            #   3. Is there a registered CommandSet that is is the only matching subclass?
+
+            # check if the command's CommandSet is a sub-class of the support function's defining class
+            if isinstance(cmd_self, func_class):
+                # Case 1: Command's CommandSet is a sub-class of the support function's CommandSet
+                func_self = cmd_self
+            else:
+                # Search all registered CommandSets
+                func_self = None
+                candidate_sets = []  # type: List[CommandSet]
+                for installed_cmd_set in self._installed_command_sets:
+                    if type(installed_cmd_set) == func_class:
+                        # Case 2: CommandSet is an exact type match for the function's CommandSet
+                        func_self = installed_cmd_set
+                        break
+
+                    # Add candidate for Case 3:
+                    if isinstance(installed_cmd_set, func_class):
+                        candidate_sets.append(installed_cmd_set)
+                if func_self is None and len(candidate_sets) == 1:
+                    # Case 3: There exists exactly 1 CommandSet that is a sub-class match of the function's CommandSet
+                    func_self = candidate_sets[0]
+            return func_self
+        else:
+            return self
diff --git a/cmd2/command_definition.py b/cmd2/command_definition.py
index 86f1151a..27a044bc 100644
--- a/cmd2/command_definition.py
+++ b/cmd2/command_definition.py
@@ -2,8 +2,7 @@
 """
 Supports the definition of commands in separate classes to be composed into cmd2.Cmd
 """
-import functools
-from typing import Callable, Iterable, Optional, Type
+from typing import Optional, Type
 
 from .constants import COMMAND_FUNC_PREFIX
 from .exceptions import CommandSetRegistrationError
diff --git a/cmd2/decorators.py b/cmd2/decorators.py
index 7c20af68..689f29c5 100644
--- a/cmd2/decorators.py
+++ b/cmd2/decorators.py
@@ -94,7 +94,7 @@ def with_argument_list(*args: List[Callable], preserve_quotes: bool = False) ->
     >>>     def do_echo(self, arglist):
     >>>         self.poutput(' '.join(arglist)
     """
-    import functools, cmd2
+    import functools
 
     def arg_decorator(func: Callable):
         @functools.wraps(func)
@@ -282,7 +282,11 @@ def with_argparser(parser: argparse.ArgumentParser, *,
             if ns_provider is None:
                 namespace = None
             else:
-                namespace = ns_provider(cmd2_app)
+                # The namespace provider may or may not be defined in the same class as the command. Since provider
+                # functions are registered with the command argparser before anything is instantiated, we
+                # need to find an instance at runtime that matches the types during declaration
+                provider_self = cmd2_app._resolve_func_self(ns_provider, args[0])
+                namespace = ns_provider(provider_self if not None else cmd2_app)
 
             try:
                 if with_unknown_args: