Refactored ArgparseCompleter to support custom completion

3 files changed, 59 insertions, 67 deletions

diff --git a/cmd2/__init__.py b/cmd2/__init__.py
index 35b163b1..0ef038f5 100644
--- a/cmd2/__init__.py
+++ b/cmd2/__init__.py
@@ -27,7 +27,7 @@ if cmd2_parser_module is not None:
 
 # Get the current value for argparse_custom.DEFAULT_ARGUMENT_PARSER
 from .argparse_custom import DEFAULT_ARGUMENT_PARSER
-from .cmd2 import Cmd
+from .cmd2 import Cmd, CompletionMode
 from .command_definition import CommandSet, with_default_category
 from .constants import COMMAND_NAME, DEFAULT_SHORTCUTS
 from .decorators import with_argument_list, with_argparser, with_argparser_and_unknown_args, with_category, \
diff --git a/cmd2/argparse_completer.py b/cmd2/argparse_completer.py
index 8bfaec80..d74226fa 100644
--- a/cmd2/argparse_completer.py
+++ b/cmd2/argparse_completer.py
@@ -182,10 +182,19 @@ class ArgparseCompleter:
                 if isinstance(action, argparse._SubParsersAction):
                     self._subcommand_action = action
 
-    def complete_command(self, tokens: List[str], text: str, line: str, begidx: int, endidx: int, *,
-                         cmd_set: Optional[CommandSet] = None) -> List[str]:
+    def complete(self, text: str, line: str, begidx: int, endidx: int, tokens: List[str], *,
+                 cmd_set: Optional[CommandSet] = None) -> List[str]:
         """
-        Complete the command using the argparse metadata and provided argument dictionary
+        Complete text using argparse metadata
+
+        :param text: the string prefix we are attempting to match (all matches must begin with it)
+        :param line: the current input line with leading whitespace removed
+        :param begidx: the beginning index of the prefix text
+        :param endidx: the ending index of the prefix text
+        :param tokens: list of argument tokens to parse
+        :param cmd_set: if tab completing a command, the CommandSet the command's function belongs to, if applicable.
+                        Defaults to None.
+
         :raises: CompletionError for various types of tab completion errors
         """
         if not tokens:
@@ -262,7 +271,7 @@ class ArgparseCompleter:
         #############################################################################################
         # Parse all but the last token
         #############################################################################################
-        for token_index, token in enumerate(tokens[1:-1], start=1):
+        for token_index, token in enumerate(tokens[:-1]):
 
             # If we're in a positional REMAINDER arg, force all future tokens to go to that
             if pos_arg_state is not None and pos_arg_state.is_remainder:
@@ -360,8 +369,8 @@ class ArgparseCompleter:
 
                                 completer = ArgparseCompleter(self._subcommand_action.choices[token], self._cmd2_app,
                                                               parent_tokens=parent_tokens)
-                                return completer.complete_command(tokens[token_index:], text, line, begidx, endidx,
-                                                                  cmd_set=cmd_set)
+                                return completer.complete(text, line, begidx, endidx, tokens[token_index + 1:],
+                                                          cmd_set=cmd_set)
                             else:
                                 # Invalid subcommand entered, so no way to complete remaining tokens
                                 return []
@@ -405,9 +414,8 @@ class ArgparseCompleter:
 
         # Check if we are completing a flag's argument
         if flag_arg_state is not None:
-            completion_results = self._complete_for_arg(flag_arg_state, text, line,
-                                                        begidx, endidx, consumed_arg_values,
-                                                        cmd_set=cmd_set)
+            completion_results = self._complete_arg(text, line, begidx, endidx, flag_arg_state, consumed_arg_values,
+                                                    cmd_set=cmd_set)
 
             # If we have results, then return them
             if completion_results:
@@ -426,9 +434,8 @@ class ArgparseCompleter:
                 action = remaining_positionals.popleft()
                 pos_arg_state = _ArgumentState(action)
 
-            completion_results = self._complete_for_arg(pos_arg_state, text, line,
-                                                        begidx, endidx, consumed_arg_values,
-                                                        cmd_set=cmd_set)
+            completion_results = self._complete_arg(text, line, begidx, endidx, pos_arg_state, consumed_arg_values,
+                                                    cmd_set=cmd_set)
 
             # If we have results, then return them
             if completion_results:
@@ -514,23 +521,23 @@ class ArgparseCompleter:
 
         return completions
 
-    def complete_subcommand_help(self, tokens: List[str], text: str, line: str, begidx: int, endidx: int) -> List[str]:
+    def complete_subcommand_help(self, text: str, line: str, begidx: int, endidx: int, tokens: List[str]) -> List[str]:
         """
         Supports cmd2's help command in the completion of subcommand names
-        :param tokens: command line tokens
         :param text: the string prefix we are attempting to match (all matches must begin with it)
         :param line: the current input line with leading whitespace removed
         :param begidx: the beginning index of the prefix text
         :param endidx: the ending index of the prefix text
+        :param tokens: arguments passed to command/subcommand
         :return: List of subcommand completions
         """
         # If our parser has subcommands, we must examine the tokens and check if they are subcommands
         # If so, we will let the subcommand's parser handle the rest of the tokens via another ArgparseCompleter.
         if self._subcommand_action is not None:
-            for token_index, token in enumerate(tokens[1:], start=1):
+            for token_index, token in enumerate(tokens):
                 if token in self._subcommand_action.choices:
                     completer = ArgparseCompleter(self._subcommand_action.choices[token], self._cmd2_app)
-                    return completer.complete_subcommand_help(tokens[token_index:], text, line, begidx, endidx)
+                    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
                     return self._cmd2_app.basic_complete(text, line, begidx, endidx, self._subcommand_action.choices)
@@ -541,24 +548,23 @@ class ArgparseCompleter:
     def format_help(self, tokens: List[str]) -> str:
         """
         Supports cmd2's help command in the retrieval of help text
-        :param tokens: command line tokens
+        :param tokens: arguments passed to help command
         :return: help text of the command being queried
         """
         # If our parser has subcommands, we must examine the tokens and check if they are subcommands
         # If so, we will let the subcommand's parser handle the rest of the tokens via another ArgparseCompleter.
         if self._subcommand_action is not None:
-            for token_index, token in enumerate(tokens[1:], start=1):
+            for token_index, token in enumerate(tokens):
                 if token in self._subcommand_action.choices:
                     completer = ArgparseCompleter(self._subcommand_action.choices[token], self._cmd2_app)
-                    return completer.format_help(tokens[token_index:])
+                    return completer.format_help(tokens[token_index + 1:])
                 else:
                     break
         return self._parser.format_help()
 
-    def _complete_for_arg(self, arg_state: _ArgumentState,
-                          text: str, line: str, begidx: int, endidx: int,
-                          consumed_arg_values: Dict[str, List[str]], *,
-                          cmd_set: Optional[CommandSet] = None) -> List[str]:
+    def _complete_arg(self, text: str, line: str, begidx: int, endidx: int,
+                      arg_state: _ArgumentState, consumed_arg_values: Dict[str, List[str]], *,
+                      cmd_set: Optional[CommandSet] = None) -> List[str]:
         """
         Tab completion routine for an argparse argument
         :return: list of completions
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 1d0f744b..c6dbef06 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -1618,7 +1618,10 @@ 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
+
         unclosed_quote = ''
+        command = None
 
         # If custom_settings is None, then we are completing a command's arguments
         if custom_settings is None:
@@ -1630,7 +1633,6 @@ class Cmd(cmd.Cmd):
             if not command:
                 return
 
-            cmd_set = self._cmd_to_command_sets[command] if command in self._cmd_to_command_sets else None
             expanded_line = statement.command_and_args
 
             # We overwrote line with a properly formatted but fully stripped version
@@ -1648,6 +1650,15 @@ class Cmd(cmd.Cmd):
             # Overwrite line to pass into completers
             line = expanded_line
 
+        # Get all tokens through the one being completed
+        tokens, raw_tokens = self.tokens_for_completion(line, begidx, endidx)
+
+        # Check if we had a parsing error
+        if len(tokens) == 0:
+            return
+
+        # Determine the completer function to use
+        if command is not None:
             # Check if a macro was entered
             if command in self.macros:
                 completer_func = self.path_complete
@@ -1663,9 +1674,12 @@ class Cmd(cmd.Cmd):
                     argparser = getattr(func, constants.CMD_ATTR_ARGPARSER, None)
 
                     if func is not None and argparser is not None:
-                        completer_func = functools.partial(self._argparse_complete,
-                                                           argparser=argparser,
-                                                           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
+                        completer = ArgparseCompleter(argparser, self)
+                        preserve_quotes = getattr(func, constants.CMD_ATTR_PRESERVE_QUOTES)
+
+                        completer_func = functools.partial(completer.complete,
+                                                           tokens=raw_tokens[1:] if preserve_quotes else tokens[1:],
                                                            cmd_set=cmd_set)
                     else:
                         completer_func = self.completedefault
@@ -1680,18 +1694,11 @@ class Cmd(cmd.Cmd):
 
         # Otherwise we are completing the command token or performing custom completion
         else:
-            completer_func = functools.partial(self._argparse_complete,
-                                               argparser=custom_settings.parser,
-                                               preserve_quotes=custom_settings.preserve_quotes,
+            completer = ArgparseCompleter(custom_settings.parser, self)
+            completer_func = functools.partial(completer.complete,
+                                               tokens=raw_tokens if custom_settings.preserve_quotes else tokens,
                                                cmd_set=None)
 
-        # Get all tokens through the one being completed
-        tokens, raw_tokens = self.tokens_for_completion(line, begidx, endidx)
-
-        # Check if we had a parsing error
-        if len(tokens) == 0:
-            return
-
         # Text we need to remove from completions later
         text_to_remove = ''
 
@@ -1873,20 +1880,6 @@ class Cmd(cmd.Cmd):
             rl_force_redisplay()
             return None
 
-    def _argparse_complete(self, text: str, line: str, begidx: int, endidx: int, *,
-                           argparser: argparse.ArgumentParser,
-                           preserve_quotes: bool,
-                           cmd_set: Optional[CommandSet] = None) -> List[str]:
-        """Perform argparse-based tab completion"""
-        from .argparse_completer import ArgparseCompleter
-        completer = ArgparseCompleter(argparser, self)
-        tokens, raw_tokens = self.tokens_for_completion(line, begidx, endidx)
-
-        # To have tab completion parsing match command line parsing behavior,
-        # use preserve_quotes to determine if we parse the quoted or unquoted tokens.
-        tokens_to_parse = raw_tokens if preserve_quotes else tokens
-        return completer.complete_command(tokens_to_parse, text, line, begidx, endidx, cmd_set=cmd_set)
-
     def in_script(self) -> bool:
         """Return whether a text script is running"""
         return self._current_script_dir is not None
@@ -2555,7 +2548,8 @@ class Cmd(cmd.Cmd):
         :param prompt: prompt to display to user
         :param history: optional list of strings to use for up-arrow history. If completion_mode is
                         CompletionMode.COMMANDS and this is None, then cmd2's command list history will
-                        be used. Defaults to None.
+                        be used. The passed in history will not be edited. It is the caller's responsibility
+                        to add the returned input to history if desired. Defaults to None.
         :param completion_mode: tells what type of tab completion to support. Tab completion only works when
                                 self.use_rawinput is True and sys.stdin is a terminal. Defaults to
                                 CompletionMode.NONE.
@@ -2565,9 +2559,9 @@ class Cmd(cmd.Cmd):
 
         A maximum of one of these should be provided:
 
-        :param choices: iterable of accepted values
-        :param choices_provider: function that provides choices for this argument
-        :param completer: tab completion function that provides choices for this argument
+        :param choices: iterable of accepted values for single argument
+        :param choices_provider: function that provides choices for single argument
+        :param completer: tab completion function that provides choices for single argument
         :param parser: an argument parser which supports the tab completion of multiple arguments
 
         :return: the line read from stdin with all trailing new lines removed
@@ -3138,12 +3132,9 @@ class Cmd(cmd.Cmd):
         if func is None or argparser is None:
             return []
 
-        # Combine the command and its subcommand tokens for the ArgparseCompleter
-        tokens = [command] + arg_tokens['subcommands']
-
         from .argparse_completer import ArgparseCompleter
         completer = ArgparseCompleter(argparser, self)
-        return completer.complete_subcommand_help(tokens, text, line, begidx, endidx)
+        return completer.complete_subcommand_help(text, line, begidx, endidx, arg_tokens['subcommands'])
 
     help_parser = DEFAULT_ARGUMENT_PARSER(description="List available commands or provide "
                                                       "detailed help for a specific command")
@@ -3174,10 +3165,9 @@ class Cmd(cmd.Cmd):
             if func is not None and argparser is not None:
                 from .argparse_completer import ArgparseCompleter
                 completer = ArgparseCompleter(argparser, self)
-                tokens = [args.command] + args.subcommands
 
                 # Set end to blank so the help output matches how it looks when "command -h" is used
-                self.poutput(completer.format_help(tokens), end='')
+                self.poutput(completer.format_help(args.subcommands), end='')
 
             # If there is no help information then print an error
             elif help_func is None and (func is None or not func.__doc__):
@@ -3382,10 +3372,6 @@ class Cmd(cmd.Cmd):
             if not response:
                 continue
 
-            if rl_type != RlType.NONE:
-                hlen = readline.get_current_history_length()
-                if hlen >= 1:
-                    readline.remove_history_item(hlen - 1)
             try:
                 choice = int(response)
                 if choice < 1:
@@ -3420,7 +3406,7 @@ class Cmd(cmd.Cmd):
 
         # Use raw_tokens since quotes have been preserved
         _, raw_tokens = self.tokens_for_completion(line, begidx, endidx)
-        return completer.complete_command(raw_tokens, text, line, begidx, endidx)
+        return completer.complete(text, line, begidx, endidx, raw_tokens[1:])
 
     # When tab completing value, we recreate the set command parser with a value argument specific to
     # the settable being edited. To make this easier, define a parent parser with all the common elements.