Merge branch 'master' into ply

2 files changed, 77 insertions, 182 deletions

diff --git a/cmd2/argparse_completer.py b/cmd2/argparse_completer.py
index 35f9342b..03f2d965 100755
--- a/cmd2/argparse_completer.py
+++ b/cmd2/argparse_completer.py
@@ -71,6 +71,10 @@ import re as _re
 from .rl_utils import rl_force_redisplay
 
 
+# attribute that can optionally added to an argparse argument (called an Action) to
+# define the completion choices for the argument. You may provide a Collection or a Function.
+ACTION_ARG_CHOICES = 'arg_choices'
+
 class _RangeAction(object):
     def __init__(self, nargs: Union[int, str, Tuple[int, int], None]):
         self.nargs_min = None
@@ -186,7 +190,8 @@ class AutoCompleter(object):
                  token_start_index: int = 1,
                  arg_choices: Dict[str, Union[List, Tuple, Callable]] = None,
                  subcmd_args_lookup: dict = None,
-                 tab_for_arg_help: bool = True):
+                 tab_for_arg_help: bool = True,
+                 cmd2_app=None):
         """
         Create an AutoCompleter
 
@@ -195,6 +200,8 @@ class AutoCompleter(object):
         :param arg_choices: dictionary mapping from argparse argument 'dest' name to list of choices
         :param subcmd_args_lookup: mapping a sub-command group name to a tuple to fill the child\
         AutoCompleter's arg_choices and subcmd_args_lookup parameters
+        :param tab_for_arg_help: Enable of disable argument help when there's no completion result
+        :param cmd2_app: reference to the Cmd2 application. Enables argparse argument completion with class methods
         """
         if not subcmd_args_lookup:
             subcmd_args_lookup = {}
@@ -205,6 +212,7 @@ class AutoCompleter(object):
         self._arg_choices = arg_choices.copy() if arg_choices is not None else {}
         self._token_start_index = token_start_index
         self._tab_for_arg_help = tab_for_arg_help
+        self._cmd2_app = cmd2_app
 
         self._flags = []  # all flags in this command
         self._flags_without_args = []  # all flags that don't take arguments
@@ -220,6 +228,10 @@ class AutoCompleter(object):
             # if there are choices defined, record them in the arguments dictionary
             if action.choices is not None:
                 self._arg_choices[action.dest] = action.choices
+            # if completion choices are tagged on the action, record them
+            elif hasattr(action, ACTION_ARG_CHOICES):
+                action_arg_choices = getattr(action, ACTION_ARG_CHOICES)
+                self._arg_choices[action.dest] = action_arg_choices
 
             # if the parameter is flag based, it will have option_strings
             if action.option_strings:
@@ -244,7 +256,8 @@ class AutoCompleter(object):
                         subcmd_start = token_start_index + len(self._positional_actions)
                         sub_completers[subcmd] = AutoCompleter(action.choices[subcmd], subcmd_start,
                                                                arg_choices=subcmd_args,
-                                                               subcmd_args_lookup=subcmd_lookup)
+                                                               subcmd_args_lookup=subcmd_lookup,
+                                                               cmd2_app=cmd2_app)
                         sub_commands.append(subcmd)
                     self._positional_completers[action.dest] = sub_completers
                     self._arg_choices[action.dest] = sub_commands
@@ -406,6 +419,21 @@ class AutoCompleter(object):
 
         return completion_results
 
+    def complete_command_help(self, tokens: List[str], text: str, line: str, begidx: int, endidx: int) -> List[str]:
+        """Supports the completion of sub-commands for commands through the cmd2 help command."""
+        for idx, token in enumerate(tokens):
+            if idx >= self._token_start_index:
+                if self._positional_completers:
+                    # For now argparse only allows 1 sub-command group per level
+                    # so this will only loop once.
+                    for completers in self._positional_completers.values():
+                        if token in completers:
+                            return completers[token].complete_command_help(tokens, text, line, begidx, endidx)
+                        else:
+                            return self.basic_complete(text, line, begidx, endidx, completers.keys())
+        return []
+
+
     @staticmethod
     def _process_action_nargs(action: argparse.Action, arg_state: _ArgumentState) -> None:
         if isinstance(action, _RangeAction):
@@ -467,8 +495,18 @@ class AutoCompleter(object):
     def _resolve_choices_for_arg(self, action: argparse.Action, used_values=()) -> List[str]:
         if action.dest in self._arg_choices:
             args = self._arg_choices[action.dest]
+
             if callable(args):
-                args = args()
+                try:
+                    if self._cmd2_app is not None:
+                        try:
+                            args = args(self._cmd2_app)
+                        except TypeError:
+                            args = args()
+                    else:
+                        args = args()
+                except TypeError:
+                    return []
 
             try:
                 iter(args)
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 7b01a653..43973b64 100755
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -50,6 +50,7 @@ import pyperclip
 
 # Set up readline
 from .rl_utils import rl_force_redisplay, readline, rl_type, RlType
+from .argparse_completer import AutoCompleter, ACArgumentParser
 
 from cmd2.parsing import CommandParser
 
@@ -261,23 +262,8 @@ def with_argparser_and_unknown_args(argparser: argparse.ArgumentParser) -> Calla
 
         cmd_wrapper.__doc__ = argparser.format_help()
 
-        # Mark this function as having an argparse ArgumentParser (used by do_help)
-        cmd_wrapper.__dict__['has_parser'] = True
-
-        # If there are subcommands, store their names in a list to support tab-completion of subcommand names
-        if argparser._subparsers is not None:
-            # Key is subcommand name and value is completer function
-            subcommands = collections.OrderedDict()
-
-            # Get all subcommands and check if they have completer functions
-            for name, parser in argparser._subparsers._group_actions[0]._name_parser_map.items():
-                if 'completer' in parser._defaults:
-                    completer = parser._defaults['completer']
-                else:
-                    completer = None
-                subcommands[name] = completer
-
-            cmd_wrapper.__dict__['subcommands'] = subcommands
+        # Mark this function as having an argparse ArgumentParser
+        setattr(cmd_wrapper, 'argparser', argparser)
 
         return cmd_wrapper
 
@@ -313,24 +299,8 @@ def with_argparser(argparser: argparse.ArgumentParser) -> Callable:
 
         cmd_wrapper.__doc__ = argparser.format_help()
 
-        # Mark this function as having an argparse ArgumentParser (used by do_help)
-        cmd_wrapper.__dict__['has_parser'] = True
-
-        # If there are subcommands, store their names in a list to support tab-completion of subcommand names
-        if argparser._subparsers is not None:
-
-            # Key is subcommand name and value is completer function
-            subcommands = collections.OrderedDict()
-
-            # Get all subcommands and check if they have completer functions
-            for name, parser in argparser._subparsers._group_actions[0]._name_parser_map.items():
-                if 'completer' in parser._defaults:
-                    completer = parser._defaults['completer']
-                else:
-                    completer = None
-                subcommands[name] = completer
-
-            cmd_wrapper.__dict__['subcommands'] = subcommands
+        # Mark this function as having an argparse ArgumentParser
+        setattr(cmd_wrapper, 'argparser', argparser)
 
         return cmd_wrapper
 
@@ -992,49 +962,6 @@ class Cmd(cmd.Cmd):
             return self._colorcodes[color][True] + val + self._colorcodes[color][False]
         return val
 
-    def get_subcommands(self, command):
-        """
-        Returns a list of a command's subcommand names if they exist
-        :param command: the command we are querying
-        :return: A subcommand list or None
-        """
-
-        subcommand_names = None
-
-        # Check if is a valid command
-        funcname = self._func_named(command)
-
-        if funcname:
-            # Check to see if this function was decorated with an argparse ArgumentParser
-            func = getattr(self, funcname)
-            subcommands = func.__dict__.get('subcommands', None)
-            if subcommands is not None:
-                subcommand_names = subcommands.keys()
-
-        return subcommand_names
-
-    def get_subcommand_completer(self, command, subcommand):
-        """
-        Returns a subcommand's tab completion function if one exists
-        :param command: command which owns the subcommand
-        :param subcommand: the subcommand we are querying
-        :return: A completer or None
-        """
-
-        completer = None
-
-        # Check if is a valid command
-        funcname = self._func_named(command)
-
-        if funcname:
-            # Check to see if this function was decorated with an argparse ArgumentParser
-            func = getattr(self, funcname)
-            subcommands = func.__dict__.get('subcommands', None)
-            if subcommands is not None:
-                completer = subcommands[subcommand]
-
-        return completer
-
     # -----  Methods related to tab completion -----
 
     def set_completion_defaults(self):
@@ -1766,16 +1693,14 @@ class Cmd(cmd.Cmd):
                     try:
                         compfunc = getattr(self, 'complete_' + command)
                     except AttributeError:
-                        compfunc = self.completedefault
-
-                    subcommands = self.get_subcommands(command)
-                    if subcommands is not None:
-                        # Since there are subcommands, then try completing those if the cursor is in
-                        # the token at index 1, otherwise default to using compfunc
-                        index_dict = {1: subcommands}
-                        compfunc = functools.partial(self.index_based_complete,
-                                                     index_dict=index_dict,
-                                                     all_else=compfunc)
+                        # There's no completer function, next see if the command uses argparser
+                        try:
+                            cmd_func = getattr(self, 'do_' + command)
+                            argparser = getattr(cmd_func, 'argparser')
+                            # Command uses argparser, switch to the default argparse completer
+                            compfunc = functools.partial(self._autocomplete_default, argparser=argparser)
+                        except AttributeError:
+                            compfunc = self.completedefault
 
                 # A valid command was not entered
                 else:
@@ -1882,6 +1807,16 @@ class Cmd(cmd.Cmd):
         except IndexError:
             return None
 
+    def _autocomplete_default(self, text: str, line: str, begidx: int, endidx: int,
+                              argparser: argparse.ArgumentParser) -> List[str]:
+        """Default completion function for argparse commands."""
+        completer = AutoCompleter(argparser, cmd2_app=self)
+
+        tokens, _ = self.tokens_for_completion(line, begidx, endidx)
+        results = completer.complete_command(tokens, text, line, begidx, endidx)
+
+        return results
+
     def get_all_commands(self):
         """
         Returns a list of all commands
@@ -1936,12 +1871,15 @@ class Cmd(cmd.Cmd):
             strs_to_match = list(topics | visible_commands)
             matches = self.basic_complete(text, line, begidx, endidx, strs_to_match)
 
-        # Check if we are completing a subcommand
-        elif index == subcmd_index:
-
-            # Match subcommands if any exist
-            command = tokens[cmd_index]
-            matches = self.basic_complete(text, line, begidx, endidx, self.get_subcommands(command))
+        # check if the command uses argparser
+        elif index >= subcmd_index:
+            try:
+                cmd_func = getattr(self, 'do_' + tokens[cmd_index])
+                parser = getattr(cmd_func, 'argparser')
+                completer = AutoCompleter(parser)
+                matches = completer.complete_command_help(tokens[1:], text, line, begidx, endidx)
+            except AttributeError:
+                pass
 
         return matches
 
@@ -2104,7 +2042,7 @@ class Cmd(cmd.Cmd):
         :param line: str - line of text read from input
         :return: bool - True if cmdloop() should exit, False otherwise
         """
-        stop = 0
+        stop = False
         try:
             statement = self._complete_statement(line)
             (stop, statement) = self.postparsing_precmd(statement)
@@ -2622,7 +2560,7 @@ Usage:  Usage: unalias [-a] name [name ...]
             if funcname:
                 # Check to see if this function was decorated with an argparse ArgumentParser
                 func = getattr(self, funcname)
-                if func.__dict__.get('has_parser', False):
+                if hasattr(func, 'argparser'):
                     # Function has an argparser, so get help based on all the arguments in case there are sub-commands
                     new_arglist = arglist[1:]
                     new_arglist.append('-h')
@@ -2845,10 +2783,10 @@ Usage:  Usage: unalias [-a] name [name ...]
         else:
             raise LookupError("Parameter '%s' not supported (type 'show' for list of parameters)." % param)
 
-    set_parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
+    set_parser = ACArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
     set_parser.add_argument('-a', '--all', action='store_true', help='display read-only settings as well')
     set_parser.add_argument('-l', '--long', action='store_true', help='describe function of parameter')
-    set_parser.add_argument('settable', nargs='*', help='[param_name] [value]')
+    set_parser.add_argument('settable', nargs=(0,2), help='[param_name] [value]')
 
     @with_argparser(set_parser)
     def do_set(self, args):
@@ -2929,87 +2867,6 @@ Usage:  Usage: unalias [-a] name [name ...]
         index_dict = {1: self.shell_cmd_complete}
         return self.index_based_complete(text, line, begidx, endidx, index_dict, self.path_complete)
 
-    def cmd_with_subs_completer(self, text, line, begidx, endidx):
-        """
-        This is a function provided for convenience to those who want an easy way to add
-        tab completion to functions that implement subcommands. By setting this as the
-        completer of the base command function, the correct completer for the chosen subcommand
-        will be called.
-
-        The use of this function requires assigning a completer function to the subcommand's parser
-        Example:
-            A command called print has a subcommands called 'names' that needs a tab completer
-            When you create the parser for names, include the completer function in the parser's defaults.
-
-            names_parser.set_defaults(func=print_names, completer=complete_print_names)
-
-            To make sure the names completer gets called, set the completer for the print function
-            in a similar fashion to what follows.
-
-            complete_print = cmd2.Cmd.cmd_with_subs_completer
-
-        When the subcommand's completer is called, this function will have stripped off all content from the
-        beginning of the command line before the subcommand, meaning the line parameter always starts with the
-        subcommand name and the index parameters reflect this change.
-
-        For instance, the command "print names -d 2" becomes "names -d 2"
-        begidx and endidx are incremented accordingly
-
-        :param text: str - the string prefix we are attempting to match (all returned matches must begin with it)
-        :param line: str - the current input line with leading whitespace removed
-        :param begidx: int - the beginning index of the prefix text
-        :param endidx: int - the ending index of the prefix text
-        :return: List[str] - a list of possible tab completions
-        """
-        # The command is the token at index 0 in the command line
-        cmd_index = 0
-
-        # The subcommand is the token at index 1 in the command line
-        subcmd_index = 1
-
-        # Get all tokens through the one being completed
-        tokens, _ = self.tokens_for_completion(line, begidx, endidx)
-        if tokens is None:
-            return []
-
-        matches = []
-
-        # Get the index of the token being completed
-        index = len(tokens) - 1
-
-        # If the token being completed is past the subcommand name, then do subcommand specific tab-completion
-        if index > subcmd_index:
-
-            # Get the command name
-            command = tokens[cmd_index]
-
-            # Get the subcommand name
-            subcommand = tokens[subcmd_index]
-
-            # Find the offset into line where the subcommand name begins
-            subcmd_start = 0
-            for cur_index in range(0, subcmd_index + 1):
-                cur_token = tokens[cur_index]
-                subcmd_start = line.find(cur_token, subcmd_start)
-
-                if cur_index != subcmd_index:
-                    subcmd_start += len(cur_token)
-
-            # Strip off everything before subcommand name
-            orig_line = line
-            line = line[subcmd_start:]
-
-            # Update the indexes
-            diff = len(orig_line) - len(line)
-            begidx -= diff
-            endidx -= diff
-
-            # Call the subcommand specific completer if it exists
-            compfunc = self.get_subcommand_completer(command, subcommand)
-            if compfunc is not None:
-                matches = compfunc(self, text, line, begidx, endidx)
-
-        return matches
 
     # noinspection PyBroadException
     def do_py(self, arg):