Added functions to enable tab completion and choices provider functions to argparse argument values

2 files changed, 177 insertions, 129 deletions

diff --git a/cmd2/argparse_completer.py b/cmd2/argparse_completer.py
index 5f4a7a87..8c539017 100644
--- a/cmd2/argparse_completer.py
+++ b/cmd2/argparse_completer.py
@@ -64,20 +64,93 @@ import sys
 
 # imports copied from argparse to support our customized argparse functions
 from argparse import ZERO_OR_MORE, ONE_OR_MORE, ArgumentError, _, _get_action_name, SUPPRESS
-from typing import List, Dict, Tuple, Callable, Union
+from typing import Any, List, Dict, Tuple, Callable, Union
 
+from . import utils
 from .ansi import ansi_aware_write, ansi_safe_wcswidth, style_error
 from .rl_utils import rl_force_redisplay
 
-from . import utils
+# Custom argparse argument attribute that means the argument's choices come from a ArgChoicesCallable
+ARG_CHOICES_CALLABLE = 'arg_choices_callable'
 
-# 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'
 ACTION_SUPPRESS_HINT = 'suppress_hint'
 ACTION_DESCRIPTIVE_COMPLETION_HEADER = 'desc_header'
 
 
+class ArgChoicesCallable:
+    """
+    Enables using a callable as the choices provider for an argparse argument.
+    While argparse has the built-in choices attribute, it is limited to an iterable.
+    """
+    def __init__(self, is_method: bool, is_completer: bool, to_call: Callable):
+        """
+        Initializer
+
+        :param is_method: True if to_call is an instance method of a cmd2 app
+        :param is_completer: True if to_call is a tab completion routine which expects
+                             the args: text, line, begidx, endidx
+        :param to_call: the callable object that will be called to provide choices for the argument
+        """
+        self.is_completer = is_completer
+        self.is_method = is_method
+        self.to_call = to_call
+
+
+def set_arg_completer_function(arg_action: argparse.Action,
+                               completer: Callable[[str, str, int, int], List[str]]):
+    """
+    Set a tab completion function for an argparse argument to provide its choices.
+
+    Note: If completer is an instance method of a cmd2 app, then use set_arg_completer_method() instead.
+
+    :param arg_action: the argument action being added to
+    :param completer: the completer function to call
+    """
+    choices_callable = ArgChoicesCallable(is_method=False, is_completer=True, to_call=completer)
+    setattr(arg_action, ARG_CHOICES_CALLABLE, choices_callable)
+
+
+def set_arg_completer_method(arg_action: argparse.Action, completer: Callable[[Any, str, str, int, int], List[str]]):
+    """
+    Set a tab completion method for an argparse argument to provide its choices.
+
+    Note: This function expects completer to be an instance method of a cmd2 app. If completer is a function,
+    then use set_arg_completer_function() instead.
+
+    :param arg_action: the argument action being added to
+    :param completer: the completer function to call
+    """
+    choices_callable = ArgChoicesCallable(is_method=True, is_completer=True, to_call=completer)
+    setattr(arg_action, ARG_CHOICES_CALLABLE, choices_callable)
+
+
+def set_arg_choices_function(arg_action: argparse.Action, choices_func: Callable[[], List[str]]):
+    """
+    Set a function for an argparse argument to provide its choices.
+
+    Note: If choices_func is an instance method of a cmd2 app, then use set_arg_choices_method() instead.
+
+    :param arg_action: the argument action being added to
+    :param choices_func: the function to call
+    """
+    choices_callable = ArgChoicesCallable(is_method=False, is_completer=False, to_call=choices_func)
+    setattr(arg_action, ARG_CHOICES_CALLABLE, choices_callable)
+
+
+def set_arg_choices_method(arg_action: argparse.Action, choices_method: Callable[[Any], List[str]]):
+    """
+    Set a method for an argparse argument to provide its choices.
+
+    Note: This function expects choices_method to be an instance method of a cmd2 app. If choices_method is a function,
+    then use set_arg_choices_function() instead.
+
+    :param arg_action: the argument action being added to
+    :param choices_method: the method to call
+    """
+    choices_callable = ArgChoicesCallable(is_method=True, is_completer=False, to_call=choices_method)
+    setattr(arg_action, ARG_CHOICES_CALLABLE, choices_callable)
+
+
 class CompletionItem(str):
     """
     Completion item with descriptive text attached
@@ -262,7 +335,7 @@ class AutoCompleter(object):
         Create an AutoCompleter
 
         :param parser: ArgumentParser instance
-        :param cmd2_app: reference to the Cmd2 application. Enables argparse argument completion with class methods
+        :param cmd2_app: reference to the Cmd2 application that owns this AutoCompleter
         :param tab_for_arg_help: Enable of disable argument help when there's no completion result
 
         # The following parameters are intended for internal use when AutoCompleter creates other AutoCompleters
@@ -300,10 +373,10 @@ class AutoCompleter(object):
             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
+            # otherwise check if a callable provides the choices for this argument
+            elif hasattr(action, ARG_CHOICES_CALLABLE):
+                arg_choice_callable = getattr(action, ARG_CHOICES_CALLABLE)
+                self._arg_choices[action.dest] = arg_choice_callable
 
             # if the parameter is flag based, it will have option_strings
             if action.option_strings:
@@ -388,9 +461,9 @@ class AutoCompleter(object):
             if not is_potential_flag(token, self._parser) and flag_action is not None:
                 flag_arg.count += 1
 
-                # does this complete a option item for the flag
+                # does this complete an option item for the flag
                 arg_choices = self._resolve_choices_for_arg(flag_action)
-                # if the current token matches the current position's autocomplete argument list,
+                # if the current token matches the current flag's autocomplete argument list,
                 # track that we've used it already.  Unless this is the current token, then keep it.
                 if not is_last_token and token in arg_choices:
                     consumed_arg_values.setdefault(flag_action.dest, [])
@@ -400,7 +473,7 @@ class AutoCompleter(object):
             """Consuming token as positional argument"""
             pos_arg.count += 1
 
-            # does this complete a option item for the flag
+            # does this complete an option item for the positional
             arg_choices = self._resolve_choices_for_arg(pos_action)
             # if the current token matches the current position's autocomplete argument list,
             # track that we've used it already.  Unless this is the current token, then keep it.
@@ -580,7 +653,6 @@ class AutoCompleter(object):
             if flag_action is not None:
                 consumed = consumed_arg_values[flag_action.dest]\
                     if flag_action.dest in consumed_arg_values else []
-                # current_items.extend(self._resolve_choices_for_arg(flag_action, consumed))
                 completion_results = self._complete_for_arg(flag_action, text, line, begidx, endidx, consumed)
                 if not completion_results:
                     self._print_action_help(flag_action)
@@ -661,73 +733,48 @@ class AutoCompleter(object):
                             return completers[token].format_help(tokens)
         return self._parser.format_help()
 
-    def _complete_for_arg(self, action: argparse.Action,
-                          text: str,
-                          line: str,
-                          begidx: int,
-                          endidx: int,
-                          used_values=()) -> List[str]:
-        if action.dest in self._arg_choices:
-            arg_choices = self._arg_choices[action.dest]
-
-            # if arg_choices is a tuple
-            #   Let's see if it's a custom completion function.  If it is, return what it provides
-            # To do this, we make sure the first element is either a callable
-            #   or it's the name of a callable in the application
-            if isinstance(arg_choices, tuple) and len(arg_choices) > 0 and \
-                    (callable(arg_choices[0]) or
-                         (isinstance(arg_choices[0], str) and hasattr(self._cmd2_app, arg_choices[0]) and
-                          callable(getattr(self._cmd2_app, arg_choices[0]))
-                          )
-                     ):
-
-                if callable(arg_choices[0]):
-                    completer = arg_choices[0]
-                else:
-                    completer = getattr(self._cmd2_app, arg_choices[0])
-
-                # extract the positional and keyword arguments from the tuple
-                list_args = None
-                kw_args = None
-                for index in range(1, len(arg_choices)):
-                    if isinstance(arg_choices[index], list) or isinstance(arg_choices[index], tuple):
-                        list_args = arg_choices[index]
-                    elif isinstance(arg_choices[index], dict):
-                        kw_args = arg_choices[index]
-
-                # call the provided function differently depending on the provided positional and keyword arguments
-                if list_args is not None and kw_args is not None:
-                    return completer(text, line, begidx, endidx, *list_args, **kw_args)
-                elif list_args is not None:
-                    return completer(text, line, begidx, endidx, *list_args)
-                elif kw_args is not None:
-                    return completer(text, line, begidx, endidx, **kw_args)
+    def _complete_for_arg(self, arg: argparse.Action,
+                          text: str, line: str, begidx: int, endidx: int, used_values=()) -> List[str]:
+        """Tab completion routine for argparse arguments"""
+
+        # Check the arg provides choices to the user
+        if arg.dest in self._arg_choices:
+            arg_choices = self._arg_choices[arg.dest]
+
+            # Check if the argument uses a specific tab completion function to provide its choices
+            if isinstance(arg_choices, ArgChoicesCallable) and arg_choices.is_completer:
+                if arg_choices.is_method:
+                    return arg_choices.to_call(self._cmd2_app, text, line, begidx, endidx)
                 else:
-                    return completer(text, line, begidx, endidx)
+                    return arg_choices.to_call(text, line, begidx, endidx)
+
+            # Otherwise use basic_complete on the choices
             else:
-                return utils.basic_complete(text, line, begidx, endidx,
-                                            self._resolve_choices_for_arg(action, used_values))
+                # Since choices can be various types like int, we must convert them to
+                # before strings before doing tab completion matching.
+                choices = [str(choice) for choice in self._resolve_choices_for_arg(arg, used_values)]
+                return utils.basic_complete(text, line, begidx, endidx, choices)
 
         return []
 
-    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]
-
-            # is the argument a string? If so, see if we can find an attribute in the
-            # application matching the string.
-            if isinstance(args, str):
-                args = getattr(self._cmd2_app, args)
+    def _resolve_choices_for_arg(self, arg: argparse.Action, used_values=()) -> List[str]:
+        """Retrieve a list of choices that are available for a particular argument"""
+        if arg.dest in self._arg_choices:
+            arg_choices = self._arg_choices[arg.dest]
 
-            # is the provided argument a callable. If so, call it
-            if callable(args):
-                try:
-                    args = args(self._cmd2_app)
-                except TypeError:
-                    args = args()
+            # Check if arg_choices is an ArgChoicesCallable that generates a choice list
+            if isinstance(arg_choices, ArgChoicesCallable):
+                if arg_choices.is_completer:
+                    # Tab completion routines are handled in other functions
+                    return []
+                else:
+                    if arg_choices.is_method:
+                        arg_choices = arg_choices.to_call(self._cmd2_app)
+                    else:
+                        arg_choices = arg_choices.to_call()
 
-            # filter out arguments we already used
-            return [arg for arg in args if arg not in used_values]
+            # Filter out arguments we already used
+            return [choice for choice in arg_choices if choice not in used_values]
 
         return []
 
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index d65b750c..da2e83b3 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -46,7 +46,7 @@ from . import ansi
 from . import constants
 from . import plugin
 from . import utils
-from .argparse_completer import AutoCompleter, ACArgumentParser, ACTION_ARG_CHOICES
+from .argparse_completer import AutoCompleter, ACArgumentParser, set_arg_choices_method, set_arg_completer_method
 from .clipboard import can_clip, get_paste_buffer, write_to_paste_buffer
 from .history import History, HistoryItem
 from .parsing import StatementParser, Statement, Macro, MacroArg, shlex_split
@@ -617,7 +617,8 @@ class Cmd(cmd.Cmd):
             if self.broken_pipe_warning:
                 sys.stderr.write(self.broken_pipe_warning)
 
-    def perror(self, msg: Any, *, end: str = '\n', apply_style: bool = True) -> None:
+    @staticmethod
+    def perror(msg: Any, *, end: str = '\n', apply_style: bool = True) -> None:
         """Print message to sys.stderr
 
         :param msg: message to print (anything convertible to a str with '{}'.format() is OK)
@@ -2380,11 +2381,11 @@ class Cmd(cmd.Cmd):
                                                       description=alias_create_description,
                                                       epilog=alias_create_epilog)
     alias_create_parser.add_argument('name', help='name of this alias')
-    setattr(alias_create_parser.add_argument('command', help='what the alias resolves to'),
-            ACTION_ARG_CHOICES, _get_commands_aliases_and_macros_for_completion)
-    setattr(alias_create_parser.add_argument('command_args', nargs=argparse.REMAINDER,
-                                             help='arguments to pass to command'),
-            ACTION_ARG_CHOICES, ('path_complete',))
+    set_arg_choices_method(alias_create_parser.add_argument('command', help='what the alias resolves to'),
+                           _get_commands_aliases_and_macros_for_completion)
+    set_arg_completer_method(alias_create_parser.add_argument('command_args', nargs=argparse.REMAINDER,
+                                                              help='arguments to pass to command'),
+                             path_complete)
     alias_create_parser.set_defaults(func=_alias_create)
 
     # alias -> delete
@@ -2392,8 +2393,8 @@ class Cmd(cmd.Cmd):
     alias_delete_description = "Delete specified aliases or all aliases if --all is used"
     alias_delete_parser = alias_subparsers.add_parser('delete', help=alias_delete_help,
                                                       description=alias_delete_description)
-    setattr(alias_delete_parser.add_argument('name', nargs='*', help='alias to delete'),
-            ACTION_ARG_CHOICES, _get_alias_names)
+    set_arg_choices_method(alias_delete_parser.add_argument('name', nargs='*', help='alias to delete'),
+                           _get_alias_names)
     alias_delete_parser.add_argument('-a', '--all', action='store_true', help="delete all aliases")
     alias_delete_parser.set_defaults(func=_alias_delete)
 
@@ -2406,8 +2407,8 @@ class Cmd(cmd.Cmd):
 
     alias_list_parser = alias_subparsers.add_parser('list', help=alias_list_help,
                                                     description=alias_list_description)
-    setattr(alias_list_parser.add_argument('name', nargs="*", help='alias to list'),
-            ACTION_ARG_CHOICES, _get_alias_names)
+    set_arg_choices_method(alias_list_parser.add_argument('name', nargs="*", help='alias to list'),
+                           _get_alias_names)
     alias_list_parser.set_defaults(func=_alias_list)
 
     # Preserve quotes since we are passing strings to other commands
@@ -2585,11 +2586,11 @@ class Cmd(cmd.Cmd):
                                                       description=macro_create_description,
                                                       epilog=macro_create_epilog)
     macro_create_parser.add_argument('name', help='name of this macro')
-    setattr(macro_create_parser.add_argument('command', help='what the macro resolves to'),
-            ACTION_ARG_CHOICES, _get_commands_aliases_and_macros_for_completion)
-    setattr(macro_create_parser.add_argument('command_args', nargs=argparse.REMAINDER,
-                                             help='arguments to pass to command'),
-            ACTION_ARG_CHOICES, ('path_complete',))
+    set_arg_choices_method(macro_create_parser.add_argument('command', help='what the macro resolves to'),
+                           _get_commands_aliases_and_macros_for_completion)
+    set_arg_completer_method(macro_create_parser.add_argument('command_args', nargs=argparse.REMAINDER,
+                                                              help='arguments to pass to command'),
+                             path_complete)
     macro_create_parser.set_defaults(func=_macro_create)
 
     # macro -> delete
@@ -2597,8 +2598,8 @@ class Cmd(cmd.Cmd):
     macro_delete_description = "Delete specified macros or all macros if --all is used"
     macro_delete_parser = macro_subparsers.add_parser('delete', help=macro_delete_help,
                                                       description=macro_delete_description)
-    setattr(macro_delete_parser.add_argument('name', nargs='*', help='macro to delete'),
-            ACTION_ARG_CHOICES, _get_macro_names)
+    set_arg_choices_method(macro_delete_parser.add_argument('name', nargs='*', help='macro to delete'),
+                           _get_macro_names)
     macro_delete_parser.add_argument('-a', '--all', action='store_true', help="delete all macros")
     macro_delete_parser.set_defaults(func=_macro_delete)
 
@@ -2610,8 +2611,8 @@ class Cmd(cmd.Cmd):
                               "Without arguments, all macros will be listed.")
 
     macro_list_parser = macro_subparsers.add_parser('list', help=macro_list_help, description=macro_list_description)
-    setattr(macro_list_parser.add_argument('name', nargs="*", help='macro to list'),
-            ACTION_ARG_CHOICES, _get_macro_names)
+    set_arg_choices_method(macro_list_parser.add_argument('name', nargs="*", help='macro to list'),
+                           _get_macro_names)
     macro_list_parser.set_defaults(func=_macro_list)
 
     # Preserve quotes since we are passing strings to other commands
@@ -2670,12 +2671,11 @@ class Cmd(cmd.Cmd):
         return matches
 
     help_parser = ACArgumentParser()
-
-    setattr(help_parser.add_argument('command', help="command to retrieve help for", nargs="?"),
-            ACTION_ARG_CHOICES, ('complete_help_command',))
-    setattr(help_parser.add_argument('subcommand', help="sub-command to retrieve help for",
-                                     nargs=argparse.REMAINDER),
-            ACTION_ARG_CHOICES, ('complete_help_subcommand',))
+    set_arg_completer_method(help_parser.add_argument('command', nargs="?", help="command to retrieve help for"),
+                             complete_help_command)
+    set_arg_completer_method(help_parser.add_argument('subcommand', nargs=argparse.REMAINDER,
+                                                      help="sub-command to retrieve help for"),
+                             complete_help_subcommand)
     help_parser.add_argument('-v', '--verbose', action='store_true',
                              help="print a list of all commands with descriptions of each")
 
@@ -2944,8 +2944,8 @@ class Cmd(cmd.Cmd):
     set_parser = ACArgumentParser(description=set_description)
     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')
-    setattr(set_parser.add_argument('param', nargs='?', help='parameter to set or view'),
-            ACTION_ARG_CHOICES, _get_settable_names)
+    set_arg_choices_method(set_parser.add_argument('param', nargs='?', help='parameter to set or view'),
+                           _get_settable_names)
     set_parser.add_argument('value', nargs='?', help='the new value for settable')
 
     @with_argparser(set_parser)
@@ -2987,11 +2987,11 @@ class Cmd(cmd.Cmd):
                 onchange_hook(old=orig_value, new=new_value)
 
     shell_parser = ACArgumentParser()
-    setattr(shell_parser.add_argument('command', help='the command to run'),
-            ACTION_ARG_CHOICES, ('shell_cmd_complete',))
-    setattr(shell_parser.add_argument('command_args', nargs=argparse.REMAINDER,
-                                      help='arguments to pass to command'),
-            ACTION_ARG_CHOICES, ('path_complete',))
+    set_arg_completer_method(shell_parser.add_argument('command', help='the command to run'),
+                             shell_cmd_complete)
+    set_arg_completer_method(shell_parser.add_argument('command_args', nargs=argparse.REMAINDER,
+                                                       help='arguments to pass to command'),
+                             path_complete)
 
     # Preserve quotes since we are passing these strings to the shell
     @with_argparser(shell_parser, preserve_quotes=True)
@@ -3052,8 +3052,8 @@ class Cmd(cmd.Cmd):
                       "by providing no arguments to py and run more complex statements there.")
 
     py_parser = ACArgumentParser(description=py_description)
-    py_parser.add_argument('command', help="command to run", nargs='?')
-    py_parser.add_argument('remainder', help="remainder of command", nargs=argparse.REMAINDER)
+    py_parser.add_argument('command', nargs='?', help="command to run")
+    py_parser.add_argument('remainder', nargs=argparse.REMAINDER, help="remainder of command")
 
     # Preserve quotes since we are passing these strings to Python
     @with_argparser(py_parser, preserve_quotes=True)
@@ -3238,11 +3238,11 @@ class Cmd(cmd.Cmd):
         return bridge.stop
 
     run_pyscript_parser = ACArgumentParser()
-    setattr(run_pyscript_parser.add_argument('script_path', help='path to the script file'),
-            ACTION_ARG_CHOICES, ('path_complete',))
-    setattr(run_pyscript_parser.add_argument('script_arguments', nargs=argparse.REMAINDER,
-                                             help='arguments to pass to script'),
-            ACTION_ARG_CHOICES, ('path_complete',))
+    set_arg_completer_method(run_pyscript_parser.add_argument('script_path', help='path to the script file'),
+                             path_complete)
+    set_arg_completer_method(run_pyscript_parser.add_argument('script_arguments', nargs=argparse.REMAINDER,
+                                                              help='arguments to pass to script'),
+                             path_complete)
 
     @with_argparser(run_pyscript_parser)
     def do_run_pyscript(self, args: argparse.Namespace) -> bool:
@@ -3300,13 +3300,13 @@ class Cmd(cmd.Cmd):
     history_action_group.add_argument('-r', '--run', action='store_true', help='run selected history items')
     history_action_group.add_argument('-e', '--edit', action='store_true',
                                       help='edit and then run selected history items')
-    setattr(history_action_group.add_argument('-o', '--output-file', metavar='FILE',
-                                              help='output commands to a script file, implies -s'),
-            ACTION_ARG_CHOICES, ('path_complete',))
-    setattr(history_action_group.add_argument('-t', '--transcript',
-                                              help='output commands and results to a transcript file,\n'
-                                                   'implies -s'),
-            ACTION_ARG_CHOICES, ('path_complete',))
+    set_arg_completer_method(history_action_group.add_argument('-o', '--output-file', metavar='FILE',
+                                                               help='output commands to a script file, implies -s'),
+                             path_complete)
+    set_arg_completer_method(history_action_group.add_argument('-t', '--transcript',
+                                                               help='output commands and results to a transcript\n'
+                                                                    'file, implies -s'),
+                             path_complete)
     history_action_group.add_argument('-c', '--clear', action='store_true', help='clear all history')
 
     history_format_group = history_parser.add_argument_group(title='formatting')
@@ -3596,8 +3596,8 @@ class Cmd(cmd.Cmd):
                         "  set editor (program-name)")
 
     edit_parser = ACArgumentParser(description=edit_description)
-    setattr(edit_parser.add_argument('file_path', help="path to a file to open in editor", nargs="?"),
-            ACTION_ARG_CHOICES, ('path_complete',))
+    set_arg_completer_method(edit_parser.add_argument('file_path', nargs="?", help="path to a file to open in editor"),
+                             path_complete)
 
     @with_argparser(edit_parser)
     def do_edit(self, args: argparse.Namespace) -> None:
@@ -3629,11 +3629,11 @@ class Cmd(cmd.Cmd):
                               )
 
     run_script_parser = ACArgumentParser(description=run_script_description)
-    setattr(run_script_parser.add_argument('-t', '--transcript',
-                                           help='record the output of the script as a transcript file'),
-            ACTION_ARG_CHOICES, ('path_complete',))
-    setattr(run_script_parser.add_argument('script_path', help="path to the script file"),
-            ACTION_ARG_CHOICES, ('path_complete',))
+    set_arg_completer_method(run_script_parser.add_argument('-t', '--transcript', help='record the output of the '
+                                                                                       'script as a transcript file'),
+                             path_complete)
+    set_arg_completer_method(run_script_parser.add_argument('script_path', help="path to the script file"),
+                             path_complete)
 
     @with_argparser(run_script_parser)
     def do_run_script(self, args: argparse.Namespace) -> Optional[bool]:
@@ -3962,7 +3962,8 @@ class Cmd(cmd.Cmd):
                 self.disable_command(cmd_name, message_to_print)
 
     # noinspection PyUnusedLocal
-    def _report_disabled_command_usage(self, *args, message_to_print: str, **kwargs) -> None:
+    @staticmethod
+    def _report_disabled_command_usage(*args, message_to_print: str, **kwargs) -> None:
         """
         Report when a disabled command has been run or had help called on it
         :param args: not used