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

6 files changed, 96 insertions, 49 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6c835259..ff37f57b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,8 @@
     * Added explicit testing against python 3.5.2 for Ubuntu 16.04, and 3.5.3 for Debian 9
     * Added fallback definition of typing.Deque (taken from 3.5.4)
     * Removed explicit type hints that fail due to a bug in 3.5.2 favoring comment-based hints instead
+    * 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
 * Other 
     * Added missing doc-string for new cmd2.Cmd __init__ parameters 
       introduced by CommandSet enhancement
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:
diff --git a/tests_isolated/test_commandset/test_commandset.py b/tests_isolated/test_commandset/test_commandset.py
index e1441fe4..bab5d536 100644
--- a/tests_isolated/test_commandset/test_commandset.py
+++ b/tests_isolated/test_commandset/test_commandset.py
@@ -261,10 +261,16 @@ class LoadableBase(cmd2.CommandSet):
     def __init__(self, dummy):
         super(LoadableBase, self).__init__()
         self._dummy = dummy  # prevents autoload
+        self._cut_called = False
 
     cut_parser = cmd2.Cmd2ArgumentParser('cut')
     cut_subparsers = cut_parser.add_subparsers(title='item', help='item to cut')
 
+    def namespace_provider(self) -> argparse.Namespace:
+        ns = argparse.Namespace()
+        ns.cut_called = self._cut_called
+        return ns
+
     @cmd2.with_argparser(cut_parser)
     def do_cut(self, ns: argparse.Namespace):
         """Cut something"""
@@ -272,6 +278,7 @@ class LoadableBase(cmd2.CommandSet):
         if handler is not None:
             # Call whatever subcommand function was selected
             handler(ns)
+            self._cut_called = True
         else:
             # No subcommand was provided, so call help
             self._cmd.pwarning('This command does nothing without sub-parsers registered')
@@ -281,9 +288,13 @@ class LoadableBase(cmd2.CommandSet):
     stir_parser = cmd2.Cmd2ArgumentParser('stir')
     stir_subparsers = stir_parser.add_subparsers(title='item', help='what to stir')
 
-    @cmd2.with_argparser(stir_parser)
+    @cmd2.with_argparser(stir_parser, ns_provider=namespace_provider)
     def do_stir(self, ns: argparse.Namespace):
         """Stir something"""
+        if not ns.cut_called:
+            self._cmd.poutput('Need to cut before stirring')
+            return
+
         handler = ns.get_handler()
         if handler is not None:
             # Call whatever subcommand function was selected
@@ -371,8 +382,8 @@ class LoadableVegetables(cmd2.CommandSet):
     bokchoy_parser.add_argument('style', completer_method=complete_style_arg)
 
     @cmd2.as_subcommand_to('cut', 'bokchoy', bokchoy_parser)
-    def cut_bokchoy(self, _: cmd2.Statement):
-        self._cmd.poutput('Bok Choy')
+    def cut_bokchoy(self, ns: argparse.Namespace):
+        self._cmd.poutput('Bok Choy: ' + ns.style)
 
 
 def test_subcommands(command_sets_manual):
@@ -498,8 +509,6 @@ def test_subcommands(command_sets_manual):
 
 def test_nested_subcommands(command_sets_manual):
     base_cmds = LoadableBase(1)
-    # fruit_cmds = LoadableFruits(1)
-    # veg_cmds = LoadableVegetables(1)
     pasta_cmds = LoadablePastaStir(1)
 
     with pytest.raises(CommandSetRegistrationError):
@@ -520,6 +529,7 @@ def test_nested_subcommands(command_sets_manual):
         stir_pasta_vigor_parser = cmd2.Cmd2ArgumentParser('vigor', add_help=False)
         stir_pasta_vigor_parser.add_argument('frequency')
 
+        # stir sauce doesn't exist anywhere, this should fail
         @cmd2.as_subcommand_to('stir sauce', 'vigorously', stir_pasta_vigor_parser)
         def stir_pasta_vigorously(self, ns: argparse.Namespace):
             self._cmd.poutput('stir the pasta vigorously')
@@ -527,6 +537,20 @@ def test_nested_subcommands(command_sets_manual):
     with pytest.raises(CommandSetRegistrationError):
         command_sets_manual.register_command_set(BadNestedSubcommands(1))
 
+    fruit_cmds = LoadableFruits(1)
+    command_sets_manual.register_command_set(fruit_cmds)
+
+    # validates custom namespace provider works correctly. Stir command will fail until
+    # the cut command is called
+    result = command_sets_manual.app_cmd('stir pasta vigorously everyminute')
+    assert 'Need to cut before stirring' in result.stdout
+
+    result = command_sets_manual.app_cmd('cut banana discs')
+    assert 'cutting banana: discs' in result.stdout
+
+    result = command_sets_manual.app_cmd('stir pasta vigorously everyminute')
+    assert 'stir the pasta vigorously' in result.stdout
+
 
 class AppWithSubCommands(cmd2.Cmd):
     """Class for testing usage of `as_subcommand_to` decorator directly in a Cmd2 subclass."""