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

1 files changed, 9 insertions, 39 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