Made CompletionError exception available to non-argparse tab completion

5 files changed, 134 insertions, 145 deletions

diff --git a/cmd2/__init__.py b/cmd2/__init__.py
index 43578e46..73d70821 100644
--- a/cmd2/__init__.py
+++ b/cmd2/__init__.py
@@ -11,7 +11,7 @@ except DistributionNotFound:
     pass
 
 from .ansi import style, fg, bg
-from .argparse_custom import Cmd2ArgumentParser, CompletionError, CompletionItem, set_default_argument_parser
+from .argparse_custom import Cmd2ArgumentParser, CompletionItem, set_default_argument_parser
 
 # Check if user has defined a module that sets a custom value for argparse_custom.DEFAULT_ARGUMENT_PARSER
 import argparse
@@ -27,4 +27,4 @@ from .constants import COMMAND_NAME, DEFAULT_SHORTCUTS
 from .decorators import categorize, with_argument_list, with_argparser, with_argparser_and_unknown_args, with_category
 from .parsing import Statement
 from .py_bridge import CommandResult
-from .utils import Settable
+from .utils import CompletionError, Settable
diff --git a/cmd2/argparse_completer.py b/cmd2/argparse_completer.py
index eae2ae28..a0c19959 100644
--- a/cmd2/argparse_completer.py
+++ b/cmd2/argparse_completer.py
@@ -16,11 +16,10 @@ from typing import Dict, List, Optional, Union
 
 from . import ansi
 from . import cmd2
-from . import utils
 from .argparse_custom import ATTR_CHOICES_CALLABLE, INFINITY, generate_range_error
 from .argparse_custom import ATTR_SUPPRESS_TAB_HINT, ATTR_DESCRIPTIVE_COMPLETION_HEADER, ATTR_NARGS_RANGE
-from .argparse_custom import ChoicesCallable, CompletionError, CompletionItem
-from .rl_utils import rl_force_redisplay
+from .argparse_custom import ChoicesCallable, CompletionItem
+from .utils import basic_complete, CompletionError
 
 # If no descriptive header is supplied, then this will be used instead
 DEFAULT_DESCRIPTIVE_HEADER = 'Description'
@@ -63,43 +62,96 @@ def _looks_like_flag(token: str, parser: argparse.ArgumentParser) -> bool:
     return True
 
 
+class _ArgumentState:
+    """Keeps state of an argument being parsed"""
+    def __init__(self, arg_action: argparse.Action) -> None:
+        self.action = arg_action
+        self.min = None
+        self.max = None
+        self.count = 0
+        self.is_remainder = (self.action.nargs == argparse.REMAINDER)
+
+        # Check if nargs is a range
+        nargs_range = getattr(self.action, ATTR_NARGS_RANGE, None)
+        if nargs_range is not None:
+            self.min = nargs_range[0]
+            self.max = nargs_range[1]
+
+        # Otherwise check against argparse types
+        elif self.action.nargs is None:
+            self.min = 1
+            self.max = 1
+        elif self.action.nargs == argparse.OPTIONAL:
+            self.min = 0
+            self.max = 1
+        elif self.action.nargs == argparse.ZERO_OR_MORE or self.action.nargs == argparse.REMAINDER:
+            self.min = 0
+            self.max = INFINITY
+        elif self.action.nargs == argparse.ONE_OR_MORE:
+            self.min = 1
+            self.max = INFINITY
+        else:
+            self.min = self.action.nargs
+            self.max = self.action.nargs
+
+
 # noinspection PyProtectedMember
-class ArgparseCompleter:
-    """Automatic command line tab completion based on argparse parameters"""
+class _ActionCompletionError(CompletionError):
+    def __init__(self, arg_action: argparse.Action, completion_error: CompletionError) -> None:
+        """
+        Adds action-specific information to a CompletionError. These are raised when
+        non-argparse related errors occur during tab completion.
+        :param arg_action: action being tab completed
+        :param completion_error: error that occurred
+        """
+        # Indent all lines of completion_error
+        indented_error = textwrap.indent(str(completion_error), '  ')
 
-    class _ArgumentState:
-        """Keeps state of an argument being parsed"""
-
-        def __init__(self, arg_action: argparse.Action) -> None:
-            self.action = arg_action
-            self.min = None
-            self.max = None
-            self.count = 0
-            self.is_remainder = (self.action.nargs == argparse.REMAINDER)
-
-            # Check if nargs is a range
-            nargs_range = getattr(self.action, ATTR_NARGS_RANGE, None)
-            if nargs_range is not None:
-                self.min = nargs_range[0]
-                self.max = nargs_range[1]
-
-            # Otherwise check against argparse types
-            elif self.action.nargs is None:
-                self.min = 1
-                self.max = 1
-            elif self.action.nargs == argparse.OPTIONAL:
-                self.min = 0
-                self.max = 1
-            elif self.action.nargs == argparse.ZERO_OR_MORE or self.action.nargs == argparse.REMAINDER:
-                self.min = 0
-                self.max = INFINITY
-            elif self.action.nargs == argparse.ONE_OR_MORE:
-                self.min = 1
-                self.max = INFINITY
-            else:
-                self.min = self.action.nargs
-                self.max = self.action.nargs
+        error = ("\nError tab completing {}:\n"
+                 "{}\n".format(argparse._get_action_name(arg_action), indented_error))
+        super().__init__(ansi.style_error(error))
+
+
+# noinspection PyProtectedMember
+class _UnfinishedFlagError(CompletionError):
+    def __init__(self, flag_arg_state: _ArgumentState) -> None:
+        """
+        CompletionError which occurs when the user has not finished the current flag
+        :param flag_arg_state: information about the unfinished flag action
+        """
+        error = "\nError: argument {}: {} ({} entered)\n".\
+            format(argparse._get_action_name(flag_arg_state.action),
+                   generate_range_error(flag_arg_state.min, flag_arg_state.max),
+                   flag_arg_state.count)
+        super().__init__(ansi.style_error(error))
+
+
+# noinspection PyProtectedMember
+class _NoResultsError(CompletionError):
+    def __init__(self, parser: argparse.ArgumentParser, arg_action: argparse.Action) -> None:
+        """
+        CompletionError which occurs when there are no results. If hinting is allowed, then its message will
+        be a hint about the argument being tab completed.
+        :param parser: ArgumentParser instance which owns the action being tab completed
+        :param arg_action: action being tab completed
+        """
+        # Check if hinting is disabled
+        suppress_hint = getattr(arg_action, ATTR_SUPPRESS_TAB_HINT, False)
+        if suppress_hint or arg_action.help == argparse.SUPPRESS:
+            hint_str = ''
+        else:
+            # Use the parser's help formatter to print just this action's help text
+            formatter = parser._get_formatter()
+            formatter.start_section("Hint")
+            formatter.add_argument(arg_action)
+            formatter.end_section()
+            hint_str = '\n' + formatter.format_help()
+        super().__init__(hint_str)
 
+
+# noinspection PyProtectedMember
+class ArgparseCompleter:
+    """Automatic command line tab completion based on argparse parameters"""
     def __init__(self, parser: argparse.ArgumentParser, cmd2_app: cmd2.Cmd, *,
                  parent_tokens: Optional[Dict[str, List[str]]] = None) -> None:
         """
@@ -141,7 +193,10 @@ class ArgparseCompleter:
                     self._subcommand_action = action
 
     def complete_command(self, tokens: List[str], text: str, line: str, begidx: int, endidx: int) -> List[str]:
-        """Complete the command using the argparse metadata and provided argument dictionary"""
+        """
+        Complete the command using the argparse metadata and provided argument dictionary
+        :raises: CompletionError for various types of tab completion errors
+        """
         if not tokens:
             return []
 
@@ -167,18 +222,18 @@ class ArgparseCompleter:
         # Completed mutually exclusive groups
         completed_mutex_groups = dict()  # dict(argparse._MutuallyExclusiveGroup -> Action which completed group)
 
-        def consume_argument(arg_state: ArgparseCompleter._ArgumentState) -> None:
+        def consume_argument(arg_state: _ArgumentState) -> None:
             """Consuming token as an argument"""
             arg_state.count += 1
             consumed_arg_values.setdefault(arg_state.action.dest, [])
             consumed_arg_values[arg_state.action.dest].append(token)
 
-        def update_mutex_groups(arg_action: argparse.Action) -> bool:
+        def update_mutex_groups(arg_action: argparse.Action) -> None:
             """
             Check if an argument belongs to a mutually exclusive group and either mark that group
             as complete or print an error if the group has already been completed
             :param arg_action: the action of the argument
-            :return: False if the group has already been completed and there is a conflict, otherwise True
+            :raises: CompletionError if the group is already completed
             """
             # Check if this action is in a mutually exclusive group
             for group in self._parser._mutually_exclusive_groups:
@@ -191,13 +246,12 @@ class ArgparseCompleter:
                         # since it's allowed to appear on the command line more than once.
                         completer_action = completed_mutex_groups[group]
                         if arg_action == completer_action:
-                            return True
+                            return
 
                         error = ansi.style_error("\nError: argument {}: not allowed with argument {}\n".
                                                  format(argparse._get_action_name(arg_action),
                                                         argparse._get_action_name(completer_action)))
-                        self._print_message(error)
-                        return False
+                        raise CompletionError(error)
 
                     # Mark that this action completed the group
                     completed_mutex_groups[group] = arg_action
@@ -214,8 +268,6 @@ class ArgparseCompleter:
                     # Arg can only be in one group, so we are done
                     break
 
-            return True
-
         #############################################################################################
         # Parse all but the last token
         #############################################################################################
@@ -238,8 +290,7 @@ class ArgparseCompleter:
             elif token == '--' and not skip_remaining_flags:
                 # Check if there is an unfinished flag
                 if flag_arg_state is not None and flag_arg_state.count < flag_arg_state.min:
-                    self._print_unfinished_flag_error(flag_arg_state)
-                    return []
+                    raise _UnfinishedFlagError(flag_arg_state)
 
                 # Otherwise end the current flag
                 else:
@@ -252,8 +303,7 @@ class ArgparseCompleter:
 
                 # Check if there is an unfinished flag
                 if flag_arg_state is not None and flag_arg_state.count < flag_arg_state.min:
-                    self._print_unfinished_flag_error(flag_arg_state)
-                    return []
+                    raise _UnfinishedFlagError(flag_arg_state)
 
                 # Reset flag arg state but not positional tracking because flags can be
                 # interspersed anywhere between positionals
@@ -269,9 +319,7 @@ class ArgparseCompleter:
                         action = self._flag_to_action[candidates_flags[0]]
 
                 if action is not None:
-                    if not update_mutex_groups(action):
-                        return []
-
+                    update_mutex_groups(action)
                     if isinstance(action, (argparse._AppendAction,
                                            argparse._AppendConstAction,
                                            argparse._CountAction)):
@@ -286,7 +334,7 @@ class ArgparseCompleter:
                         # earlier in the command line. Reset them now for this use of it.
                         consumed_arg_values[action.dest] = []
 
-                    new_arg_state = ArgparseCompleter._ArgumentState(action)
+                    new_arg_state = _ArgumentState(action)
 
                     # Keep track of this flag if it can receive arguments
                     if new_arg_state.max > 0:
@@ -328,14 +376,11 @@ class ArgparseCompleter:
 
                         # Otherwise keep track of the argument
                         else:
-                            pos_arg_state = ArgparseCompleter._ArgumentState(action)
+                            pos_arg_state = _ArgumentState(action)
 
                 # Check if we have a positional to consume this token
                 if pos_arg_state is not None:
-                    # No need to check for an error since we remove a completed group's positional from
-                    # remaining_positionals which means this action can't belong to a completed mutex group
                     update_mutex_groups(pos_arg_state.action)
-
                     consume_argument(pos_arg_state)
 
                     # No more flags are allowed if this is a REMAINDER argument
@@ -361,9 +406,7 @@ class ArgparseCompleter:
         # character (-f) at the end.
         if _looks_like_flag(text, self._parser) and not skip_remaining_flags:
             if flag_arg_state is not None and flag_arg_state.count < flag_arg_state.min:
-                self._print_unfinished_flag_error(flag_arg_state)
-                return []
-
+                raise _UnfinishedFlagError(flag_arg_state)
             return self._complete_flags(text, line, begidx, endidx, matched_flags)
 
         completion_results = []
@@ -374,8 +417,7 @@ class ArgparseCompleter:
                 completion_results = self._complete_for_arg(flag_arg_state.action, text, line,
                                                             begidx, endidx, consumed_arg_values)
             except CompletionError as ex:
-                self._print_completion_error(flag_arg_state.action, ex)
-                return []
+                raise _ActionCompletionError(flag_arg_state.action, ex)
 
             # If we have results, then return them
             if completion_results:
@@ -384,8 +426,7 @@ class ArgparseCompleter:
             # Otherwise, print a hint if the flag isn't finished or text isn't possibly the start of a flag
             elif flag_arg_state.count < flag_arg_state.min or \
                     not _single_prefix_char(text, self._parser) or skip_remaining_flags:
-                self._print_arg_hint(flag_arg_state.action)
-                return []
+                raise _NoResultsError(self._parser, flag_arg_state.action)
 
         # Otherwise check if we have a positional to complete
         elif pos_arg_state is not None or remaining_positionals:
@@ -393,14 +434,13 @@ class ArgparseCompleter:
             # If we aren't current tracking a positional, then get the next positional arg to handle this token
             if pos_arg_state is None:
                 action = remaining_positionals.popleft()
-                pos_arg_state = ArgparseCompleter._ArgumentState(action)
+                pos_arg_state = _ArgumentState(action)
 
             try:
                 completion_results = self._complete_for_arg(pos_arg_state.action, text, line,
                                                             begidx, endidx, consumed_arg_values)
             except CompletionError as ex:
-                self._print_completion_error(pos_arg_state.action, ex)
-                return []
+                raise _ActionCompletionError(pos_arg_state.action, ex)
 
             # If we have results, then return them
             if completion_results:
@@ -408,8 +448,7 @@ class ArgparseCompleter:
 
             # Otherwise, print a hint if text isn't possibly the start of a flag
             elif not _single_prefix_char(text, self._parser) or skip_remaining_flags:
-                self._print_arg_hint(pos_arg_state.action)
-                return []
+                raise _NoResultsError(self._parser, pos_arg_state.action)
 
         # Handle case in which text is a single flag prefix character that
         # didn't complete against any argument values.
@@ -432,7 +471,7 @@ class ArgparseCompleter:
                 if action.help != argparse.SUPPRESS:
                     match_against.append(flag)
 
-        return utils.basic_complete(text, line, begidx, endidx, match_against)
+        return basic_complete(text, line, begidx, endidx, match_against)
 
     def _format_completions(self, action, completions: List[Union[str, CompletionItem]]) -> List[str]:
         # Check if the results are CompletionItems and that there aren't too many to display
@@ -491,7 +530,7 @@ class ArgparseCompleter:
                     return completer.complete_subcommand_help(tokens[token_index:], text, line, begidx, endidx)
                 elif token_index == len(tokens) - 1:
                     # Since this is the last token, we will attempt to complete it
-                    return utils.basic_complete(text, line, begidx, endidx, self._subcommand_action.choices)
+                    return basic_complete(text, line, begidx, endidx, self._subcommand_action.choices)
                 else:
                     break
         return []
@@ -519,7 +558,7 @@ class ArgparseCompleter:
         """
         Tab completion routine for an argparse argument
         :return: list of completions
-        :raises CompletionError if the completer or choices function this calls raises one
+        :raises: CompletionError if the completer or choices function this calls raises one
         """
         # Check if the arg provides choices to the user
         if arg_action.choices is not None:
@@ -579,55 +618,6 @@ class ArgparseCompleter:
             arg_choices = [choice for choice in arg_choices if choice not in used_values]
 
             # Do tab completion on the choices
-            results = utils.basic_complete(text, line, begidx, endidx, arg_choices)
+            results = basic_complete(text, line, begidx, endidx, arg_choices)
 
         return self._format_completions(arg_action, results)
-
-    @staticmethod
-    def _print_message(msg: str) -> None:
-        """Print a message instead of tab completions and redraw the prompt and input line"""
-        import sys
-        ansi.style_aware_write(sys.stdout, msg + '\n')
-        rl_force_redisplay()
-
-    def _print_arg_hint(self, arg_action: argparse.Action) -> None:
-        """
-        Print argument hint to the terminal when tab completion results in no results
-        :param arg_action: action being tab completed
-        """
-        # Check if hinting is disabled
-        suppress_hint = getattr(arg_action, ATTR_SUPPRESS_TAB_HINT, False)
-        if suppress_hint or arg_action.help == argparse.SUPPRESS:
-            return
-
-        # Use the parser's help formatter to print just this action's help text
-        formatter = self._parser._get_formatter()
-        formatter.start_section("Hint")
-        formatter.add_argument(arg_action)
-        formatter.end_section()
-        out_str = formatter.format_help()
-        self._print_message('\n' + out_str)
-
-    def _print_unfinished_flag_error(self, flag_arg_state: _ArgumentState) -> None:
-        """
-        Print an error during tab completion when the user has not finished the current flag
-        :param flag_arg_state: information about the unfinished flag action
-        """
-        error = "\nError: argument {}: {} ({} entered)\n".\
-            format(argparse._get_action_name(flag_arg_state.action),
-                   generate_range_error(flag_arg_state.min, flag_arg_state.max),
-                   flag_arg_state.count)
-        self._print_message(ansi.style_error('{}'.format(error)))
-
-    def _print_completion_error(self, arg_action: argparse.Action, completion_error: CompletionError) -> None:
-        """
-        Print a CompletionError to the user
-        :param arg_action: action being tab completed
-        :param completion_error: error that occurred
-        """
-        # Indent all lines of completion_error
-        indented_error = textwrap.indent(str(completion_error), '  ')
-
-        error = ("\nError tab completing {}:\n"
-                 "{}\n".format(argparse._get_action_name(arg_action), indented_error))
-        self._print_message(ansi.style_error('{}'.format(error)))
diff --git a/cmd2/argparse_custom.py b/cmd2/argparse_custom.py
index fd1ea057..81fec013 100644
--- a/cmd2/argparse_custom.py
+++ b/cmd2/argparse_custom.py
@@ -117,13 +117,6 @@ Tab Completion:
     their values. All tokens are stored in the dictionary as the raw strings provided on the command line. It is up to
     the developer to determine if the user entered the correct argument type (e.g. int) and validate their values.
 
-CompletionError Class:
-    Raised during tab completion operations to report any sort of error you want printed by the ArgparseCompleter
-
-    Example use cases
-    - Reading a database to retrieve a tab completion data set failed
-    - A previous command line argument that determines the data set being completed is invalid
-
 CompletionItem Class:
     This class was added to help in cases where uninformative data is being tab completed. For instance,
     tab completing ID numbers isn't very helpful to a user without context. Returning a list of CompletionItems
@@ -229,17 +222,6 @@ def generate_range_error(range_min: int, range_max: Union[int, float]) -> str:
     return err_str
 
 
-class CompletionError(Exception):
-    """
-    Raised during tab completion operations to report any sort of error you want printed by the ArgparseCompleter
-
-    Example use cases
-    - Reading a database to retrieve a tab completion data set failed
-    - A previous command line argument that determines the data set being completed is invalid
-    """
-    pass
-
-
 class CompletionItem(str):
     """
     Completion item with descriptive text attached
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 9e1085b2..f4e1ef8d 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -47,13 +47,13 @@ from . import ansi
 from . import constants
 from . import plugin
 from . import utils
-from .argparse_custom import CompletionError, CompletionItem, DEFAULT_ARGUMENT_PARSER
+from .argparse_custom import CompletionItem, DEFAULT_ARGUMENT_PARSER
 from .clipboard import can_clip, get_paste_buffer, write_to_paste_buffer
 from .decorators import with_argparser
 from .history import History, HistoryItem
 from .parsing import StatementParser, Statement, Macro, MacroArg, shlex_split
 from .rl_utils import rl_type, RlType, rl_get_point, rl_set_prompt, vt100_support, rl_make_safe_prompt, rl_warning
-from .utils import Settable
+from .utils import CompletionError, Settable
 
 # Set up readline
 if rl_type == RlType.NONE:  # pragma: no cover
@@ -1416,6 +1416,12 @@ class Cmd(cmd.Cmd):
             except IndexError:
                 return None
 
+        except CompletionError as e:
+            err_str = str(e)
+            if err_str:
+                ansi.style_aware_write(sys.stdout, err_str + '\n')
+                rl_force_redisplay()
+            return None
         except Exception as e:
             # Insert a newline so the exception doesn't print in the middle of the command line being tab completed
             self.perror()
diff --git a/cmd2/utils.py b/cmd2/utils.py
index b307e0d2..b6b45891 100644
--- a/cmd2/utils.py
+++ b/cmd2/utils.py
@@ -72,6 +72,17 @@ def str_to_bool(val: str) -> bool:
     raise ValueError("must be True or False (case-insensitive)")
 
 
+class CompletionError(Exception):
+    """
+    Raised during tab completion operations to report any sort of error you want printed by the ArgparseCompleter
+
+    Example use cases
+    - Reading a database to retrieve a tab completion data set failed
+    - A previous command line argument that determines the data set being completed is invalid
+    """
+    pass
+
+
 class Settable:
     """Used to configure a cmd2 instance member to be settable via the set command in the CLI"""
     def __init__(self, name: str, val_type: Callable, description: str, *,