Patched argparse._ActionsContainer.add_argument() to support more settings like enabling tab completion and providing choice generating functions

2 files changed, 97 insertions, 109 deletions

diff --git a/cmd2/argparse_completer.py b/cmd2/argparse_completer.py
index 8c539017..e0e38a1d 100644
--- a/cmd2/argparse_completer.py
+++ b/cmd2/argparse_completer.py
@@ -64,17 +64,20 @@ 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 Any, List, Dict, Tuple, Callable, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from . import utils
 from .ansi import ansi_aware_write, ansi_safe_wcswidth, style_error
 from .rl_utils import rl_force_redisplay
 
-# Custom argparse argument attribute that means the argument's choices come from a ArgChoicesCallable
+# Argparse argument attribute that stores an ArgChoicesCallable
 ARG_CHOICES_CALLABLE = 'arg_choices_callable'
 
-ACTION_SUPPRESS_HINT = 'suppress_hint'
-ACTION_DESCRIPTIVE_COMPLETION_HEADER = 'desc_header'
+# Argparse argument attribute that suppresses tab-completion hints
+ARG_SUPPRESS_HINT = 'arg_suppress_hint'
+
+# Argparse argument attribute that prints descriptive header when using CompletionItems
+ARG_DESCRIPTIVE_COMPLETION_HEADER = 'desc_header'
 
 
 class ArgChoicesCallable:
@@ -85,70 +88,75 @@ class ArgChoicesCallable:
     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_method: True if to_call is an instance method of a cmd2 app. False if it is a function.
         :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.is_completer = is_completer
         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.
+# Save the original _ActionsContainer.add_argument because we need to patch it
+actual_actions_container_add_argument = argparse._ActionsContainer.add_argument
 
-    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
+def patched_add_argument(self, *args,
+                         choices_function: Optional[Callable[[], List[str]]] = None,
+                         choices_method: Optional[Callable[[Any], List[str]]] = None,
+                         completer_function: Optional[Callable[[str, str, int, int], List[str]]] = None,
+                         completer_method: Optional[Callable[[Any, str, str, int, int], List[str]]] = None,
+                         suppress_hint: bool = False,
+                         description_header: Optional[str] = None,
+                         **kwargs):
     """
-    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]]):
+    This is a patched version of _ActionsContainer.add_argument() that supports more settings needed by cmd2
+    :param self:
+    :param args:
+    :param choices_function:
+    :param choices_method:
+    :param completer_function:
+    :param completer_method:
+    :param suppress_hint:
+    :param description_header:
+    :param kwargs:
+    :return:
     """
-    Set a tab completion method for an argparse argument to provide its choices.
+    # Call the actual add_argument function
+    new_arg = actual_actions_container_add_argument(self, *args, **kwargs)
 
-    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.
+    # Verify consistent use of arguments
+    choice_params = [new_arg.choices, choices_function, choices_method, completer_function, completer_method]
+    num_set = len(choice_params) - choice_params.count(None)
 
-    :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)
+    if num_set > 1:
+        err_msg = ("Only one of the following may be used in an argparser argument at a time:\n"
+                   "choices, choices_function, choices_method, completer_function, completer_method")
+        raise (ValueError(err_msg))
 
+    # Set the custom attributes
+    if choices_function:
+        setattr(new_arg, ARG_CHOICES_CALLABLE,
+                ArgChoicesCallable(is_method=False, is_completer=False, to_call=choices_function))
+    elif choices_method:
+        setattr(new_arg, ARG_CHOICES_CALLABLE,
+                ArgChoicesCallable(is_method=True, is_completer=False, to_call=choices_method))
+    elif completer_function:
+        setattr(new_arg, ARG_CHOICES_CALLABLE,
+                ArgChoicesCallable(is_method=False, is_completer=True, to_call=completer_function))
+    elif completer_method:
+        setattr(new_arg, ARG_CHOICES_CALLABLE,
+                ArgChoicesCallable(is_method=True, is_completer=True, to_call=completer_method))
 
-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.
+    setattr(new_arg, ARG_SUPPRESS_HINT, suppress_hint)
+    setattr(new_arg, ARG_DESCRIPTIVE_COMPLETION_HEADER, description_header)
 
-    Note: If choices_func is an instance method of a cmd2 app, then use set_arg_choices_method() instead.
+    return new_arg
 
-    :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)
+# Overwrite _ActionsContainer.add_argument with our patched version
+argparse._ActionsContainer.add_argument = patched_add_argument
 
 
 class CompletionItem(str):
@@ -697,7 +705,7 @@ class AutoCompleter(object):
                 completions_with_desc.append(entry)
 
             try:
-                desc_header = action.desc_header
+                desc_header = getattr(action, ARG_DESCRIPTIVE_COMPLETION_HEADER)
             except AttributeError:
                 desc_header = 'Description'
             header = '\n{: <{token_width}}{}'.format(action.dest.upper(), desc_header, token_width=token_width + 2)
@@ -784,13 +792,9 @@ class AutoCompleter(object):
             return
 
         # is parameter hinting disabled for this parameter?
-        try:
-            suppress_hint = getattr(action, ACTION_SUPPRESS_HINT)
-        except AttributeError:
-            pass
-        else:
-            if suppress_hint:
-                return
+        suppress_hint = getattr(action, ARG_SUPPRESS_HINT, False)
+        if suppress_hint:
+            return
 
         if action.option_strings:
             flags = ', '.join(action.option_strings)
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index da2e83b3..0272a7c4 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, set_arg_choices_method, set_arg_completer_method
+from .argparse_completer import AutoCompleter, ACArgumentParser
 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
@@ -2381,11 +2381,10 @@ class Cmd(cmd.Cmd):
                                                       description=alias_create_description,
                                                       epilog=alias_create_epilog)
     alias_create_parser.add_argument('name', help='name of this alias')
-    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.add_argument('command', help='what the alias resolves to',
+                                     choices_method=_get_commands_aliases_and_macros_for_completion)
+    alias_create_parser.add_argument('command_args', nargs=argparse.REMAINDER, help='arguments to pass to command',
+                                     completer_method=path_complete)
     alias_create_parser.set_defaults(func=_alias_create)
 
     # alias -> delete
@@ -2393,8 +2392,7 @@ 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)
-    set_arg_choices_method(alias_delete_parser.add_argument('name', nargs='*', help='alias to delete'),
-                           _get_alias_names)
+    alias_delete_parser.add_argument('name', nargs='*', help='alias to delete', choices_method=_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)
 
@@ -2407,8 +2405,7 @@ class Cmd(cmd.Cmd):
 
     alias_list_parser = alias_subparsers.add_parser('list', help=alias_list_help,
                                                     description=alias_list_description)
-    set_arg_choices_method(alias_list_parser.add_argument('name', nargs="*", help='alias to list'),
-                           _get_alias_names)
+    alias_list_parser.add_argument('name', nargs="*", help='alias to list', choices_method=_get_alias_names)
     alias_list_parser.set_defaults(func=_alias_list)
 
     # Preserve quotes since we are passing strings to other commands
@@ -2586,11 +2583,10 @@ class Cmd(cmd.Cmd):
                                                       description=macro_create_description,
                                                       epilog=macro_create_epilog)
     macro_create_parser.add_argument('name', help='name of this macro')
-    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.add_argument('command', help='what the macro resolves to',
+                                     choices_method=_get_commands_aliases_and_macros_for_completion)
+    macro_create_parser.add_argument('command_args', nargs=argparse.REMAINDER,
+                                     help='arguments to pass to command', completer_method=path_complete)
     macro_create_parser.set_defaults(func=_macro_create)
 
     # macro -> delete
@@ -2598,8 +2594,7 @@ 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)
-    set_arg_choices_method(macro_delete_parser.add_argument('name', nargs='*', help='macro to delete'),
-                           _get_macro_names)
+    macro_delete_parser.add_argument('name', nargs='*', help='macro to delete', choices_method=_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)
 
@@ -2611,8 +2606,7 @@ 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)
-    set_arg_choices_method(macro_list_parser.add_argument('name', nargs="*", help='macro to list'),
-                           _get_macro_names)
+    macro_list_parser.add_argument('name', nargs="*", help='macro to list', choices_method=_get_macro_names)
     macro_list_parser.set_defaults(func=_macro_list)
 
     # Preserve quotes since we are passing strings to other commands
@@ -2671,11 +2665,10 @@ class Cmd(cmd.Cmd):
         return matches
 
     help_parser = ACArgumentParser()
-    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('command', nargs="?", help="command to retrieve help for",
+                             completer_method=complete_help_command)
+    help_parser.add_argument('subcommand', nargs=argparse.REMAINDER, help="sub-command to retrieve help for",
+                             completer_method=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 +2937,7 @@ 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')
-    set_arg_choices_method(set_parser.add_argument('param', nargs='?', help='parameter to set or view'),
-                           _get_settable_names)
+    set_parser.add_argument('param', nargs='?', help='parameter to set or view', choices_method=_get_settable_names)
     set_parser.add_argument('value', nargs='?', help='the new value for settable')
 
     @with_argparser(set_parser)
@@ -2987,11 +2979,9 @@ class Cmd(cmd.Cmd):
                 onchange_hook(old=orig_value, new=new_value)
 
     shell_parser = ACArgumentParser()
-    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)
+    shell_parser.add_argument('command', help='the command to run', completer_method=shell_cmd_complete)
+    shell_parser.add_argument('command_args', nargs=argparse.REMAINDER, help='arguments to pass to command',
+                              completer_method=path_complete)
 
     # Preserve quotes since we are passing these strings to the shell
     @with_argparser(shell_parser, preserve_quotes=True)
@@ -3238,11 +3228,9 @@ class Cmd(cmd.Cmd):
         return bridge.stop
 
     run_pyscript_parser = ACArgumentParser()
-    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)
+    run_pyscript_parser.add_argument('script_path', help='path to the script file', completer_method=path_complete)
+    run_pyscript_parser.add_argument('script_arguments', nargs=argparse.REMAINDER,
+                                     help='arguments to pass to script', completer_method=path_complete)
 
     @with_argparser(run_pyscript_parser)
     def do_run_pyscript(self, args: argparse.Namespace) -> bool:
@@ -3300,13 +3288,12 @@ 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')
-    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('-o', '--output-file', metavar='FILE',
+                                      help='output commands to a script file, implies -s',
+                                      completer_method=path_complete)
+    history_action_group.add_argument('-t', '--transcript', help='output commands and results to a transcript file,\n'
+                                                                 'implies -s',
+                                      completer_method=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 +3583,8 @@ class Cmd(cmd.Cmd):
                         "  set editor (program-name)")
 
     edit_parser = ACArgumentParser(description=edit_description)
-    set_arg_completer_method(edit_parser.add_argument('file_path', nargs="?", help="path to a file to open in editor"),
-                             path_complete)
+    edit_parser.add_argument('file_path', nargs="?",
+                             help="path to a file to open in editor", completer_method=path_complete)
 
     @with_argparser(edit_parser)
     def do_edit(self, args: argparse.Namespace) -> None:
@@ -3625,15 +3612,12 @@ class Cmd(cmd.Cmd):
                               "typed in the console.\n"
                               "\n"
                               "If the -r/--record_transcript flag is used, this command instead records\n"
-                              "the output of the script commands to a transcript for testing purposes.\n"
-                              )
+                              "the output of the script commands to a transcript for testing purposes.\n")
 
     run_script_parser = ACArgumentParser(description=run_script_description)
-    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)
+    run_script_parser.add_argument('-t', '--transcript', help='record the output of the script as a transcript file',
+                                   completer_method=path_complete)
+    run_script_parser.add_argument('script_path', help="path to the script file", completer_method=path_complete)
 
     @with_argparser(run_script_parser)
     def do_run_script(self, args: argparse.Namespace) -> Optional[bool]: