Merge pull request #975 from python-cmd2/help_format_fix

Fix tuple metavar crash in Cmd2HelpFormatter and ArgparseCompleter

5 files changed, 133 insertions, 38 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ff37f57b..7361e195 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,8 @@
     * Renamed `install_command_set()` and `uninstall_command_set()` to `register_command_set()` and 
       `unregister_command_set()` for better name consistency.
 * Bug Fixes
+    * Fixed help formatting bug in `Cmd2ArgumentParser` when `metavar` is a tuple
+    * Fixed tab completion bug when using `CompletionItem` on an argument whose `metavar` is a tuple
     * Added explicit testing against python 3.5.2 for Ubuntu 16.04, and 3.5.3 for Debian 9
     * Added fallback definition of typing.Deque (taken from 3.5.4)
     * Removed explicit type hints that fail due to a bug in 3.5.2 favoring comment-based hints instead
diff --git a/cmd2/argparse_completer.py b/cmd2/argparse_completer.py
index 582f57f6..0efaebe9 100644
--- a/cmd2/argparse_completer.py
+++ b/cmd2/argparse_completer.py
@@ -25,7 +25,7 @@ from .argparse_custom import (
 )
 from .command_definition import CommandSet
 from .table_creator import Column, SimpleTable
-from .utils import CompletionError, basic_complete, get_defining_class
+from .utils import CompletionError, basic_complete
 
 # If no descriptive header is supplied, then this will be used instead
 DEFAULT_DESCRIPTIVE_HEADER = 'Description'
@@ -405,7 +405,7 @@ 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.action, text, line,
+            completion_results = self._complete_for_arg(flag_arg_state, text, line,
                                                         begidx, endidx, consumed_arg_values,
                                                         cmd_set=cmd_set)
 
@@ -426,7 +426,7 @@ class ArgparseCompleter:
                 action = remaining_positionals.popleft()
                 pos_arg_state = _ArgumentState(action)
 
-            completion_results = self._complete_for_arg(pos_arg_state.action, text, line,
+            completion_results = self._complete_for_arg(pos_arg_state, text, line,
                                                         begidx, endidx, consumed_arg_values,
                                                         cmd_set=cmd_set)
 
@@ -461,7 +461,7 @@ class ArgparseCompleter:
 
         return basic_complete(text, line, begidx, endidx, match_against)
 
-    def _format_completions(self, action, completions: List[Union[str, CompletionItem]]) -> List[str]:
+    def _format_completions(self, arg_state: _ArgumentState, completions: List[Union[str, CompletionItem]]) -> List[str]:
         # Check if the results are CompletionItems and that there aren't too many to display
         if 1 < len(completions) <= self._cmd2_app.max_completion_items and \
                 isinstance(completions[0], CompletionItem):
@@ -472,9 +472,18 @@ class ArgparseCompleter:
                 self._cmd2_app.matches_sorted = True
 
             # If a metavar was defined, use that instead of the dest field
-            destination = action.metavar if action.metavar else action.dest
-
-            desc_header = getattr(action, ATTR_DESCRIPTIVE_COMPLETION_HEADER, None)
+            destination = arg_state.action.metavar if arg_state.action.metavar else arg_state.action.dest
+
+            # Handle case where metavar was a tuple
+            if isinstance(destination, tuple):
+                # Figure out what string in the tuple to use based on how many of the arguments have been completed.
+                # Use min() to avoid going passed the end of the tuple to support nargs being ZERO_OR_MORE and
+                # ONE_OR_MORE. In those cases, argparse limits metavar tuple to 2 elements but we may be completing
+                # the 3rd or more argument here.
+                tuple_index = min(len(destination) - 1, arg_state.count)
+                destination = destination[tuple_index]
+
+            desc_header = getattr(arg_state.action, ATTR_DESCRIPTIVE_COMPLETION_HEADER, None)
             if desc_header is None:
                 desc_header = DEFAULT_DESCRIPTIVE_HEADER
 
@@ -546,7 +555,7 @@ class ArgparseCompleter:
                     break
         return self._parser.format_help()
 
-    def _complete_for_arg(self, arg_action: argparse.Action,
+    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]:
@@ -556,10 +565,10 @@ class ArgparseCompleter:
         :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:
-            arg_choices = arg_action.choices
+        if arg_state.action.choices is not None:
+            arg_choices = arg_state.action.choices
         else:
-            arg_choices = getattr(arg_action, ATTR_CHOICES_CALLABLE, None)
+            arg_choices = getattr(arg_state.action, ATTR_CHOICES_CALLABLE, None)
 
         if arg_choices is None:
             return []
@@ -586,8 +595,8 @@ class ArgparseCompleter:
                 arg_tokens = {**self._parent_tokens, **consumed_arg_values}
 
                 # Include the token being completed
-                arg_tokens.setdefault(arg_action.dest, [])
-                arg_tokens[arg_action.dest].append(text)
+                arg_tokens.setdefault(arg_state.action.dest, [])
+                arg_tokens[arg_state.action.dest].append(text)
 
                 # Add the namespace to the keyword arguments for the function we are calling
                 kwargs[ARG_TOKENS] = arg_tokens
@@ -617,10 +626,10 @@ class ArgparseCompleter:
                     arg_choices[index] = str(choice)
 
             # Filter out arguments we already used
-            used_values = consumed_arg_values.get(arg_action.dest, [])
+            used_values = consumed_arg_values.get(arg_state.action.dest, [])
             arg_choices = [choice for choice in arg_choices if choice not in used_values]
 
             # Do tab completion on the choices
             results = basic_complete(text, line, begidx, endidx, arg_choices)
 
-        return self._format_completions(arg_action, results)
+        return self._format_completions(arg_state, results)
diff --git a/cmd2/argparse_custom.py b/cmd2/argparse_custom.py
index d724cb88..12c18644 100644
--- a/cmd2/argparse_custom.py
+++ b/cmd2/argparse_custom.py
@@ -733,7 +733,8 @@ class Cmd2HelpFormatter(argparse.RawTextHelpFormatter):
                 return ', '.join(action.option_strings) + ' ' + args_string
             # End cmd2 customization
 
-    def _metavar_formatter(self, action, default_metavar) -> Callable:
+    def _determine_metavar(self, action, default_metavar) -> Union[str, Tuple]:
+        """Custom method to determine what to use as the metavar value of an action"""
         if action.metavar is not None:
             result = action.metavar
         elif action.choices is not None:
@@ -743,38 +744,46 @@ class Cmd2HelpFormatter(argparse.RawTextHelpFormatter):
             # End cmd2 customization
         else:
             result = default_metavar
+        return result
+
+    def _metavar_formatter(self, action, default_metavar) -> Callable:
+        metavar = self._determine_metavar(action, default_metavar)
 
         # noinspection PyMissingOrEmptyDocstring
         def format(tuple_size):
-            if isinstance(result, tuple):
-                return result
+            if isinstance(metavar, tuple):
+                return metavar
             else:
-                return (result, ) * tuple_size
+                return (metavar, ) * tuple_size
         return format
 
     # noinspection PyProtectedMember
     def _format_args(self, action, default_metavar) -> str:
-        get_metavar = self._metavar_formatter(action, default_metavar)
-        # Begin cmd2 customization (less verbose)
-        nargs_range = getattr(action, ATTR_NARGS_RANGE, None)
+        """Customized to handle ranged nargs and make other output less verbose"""
+        metavar = self._determine_metavar(action, default_metavar)
+        metavar_formatter = self._metavar_formatter(action, default_metavar)
 
+        # Handle nargs specified as a range
+        nargs_range = getattr(action, ATTR_NARGS_RANGE, None)
         if nargs_range is not None:
             if nargs_range[1] == constants.INFINITY:
                 range_str = '{}+'.format(nargs_range[0])
             else:
                 range_str = '{}..{}'.format(nargs_range[0], nargs_range[1])
 
-            result = '{}{{{}}}'.format('%s' % get_metavar(1), range_str)
-        elif action.nargs == ZERO_OR_MORE:
-            result = '[%s [...]]' % get_metavar(1)
-        elif action.nargs == ONE_OR_MORE:
-            result = '%s [...]' % get_metavar(1)
-        elif isinstance(action.nargs, int) and action.nargs > 1:
-            result = '{}{{{}}}'.format('%s' % get_metavar(1), action.nargs)
-        # End cmd2 customization
-        else:
-            result = super()._format_args(action, default_metavar)
-        return result
+            return '{}{{{}}}'.format('%s' % metavar_formatter(1), range_str)
+
+        # Make this output less verbose. Do not customize the output when metavar is a
+        # tuple of strings. Allow argparse's formatter to handle that instead.
+        elif isinstance(metavar, str):
+            if action.nargs == ZERO_OR_MORE:
+                return '[%s [...]]' % metavar_formatter(1)
+            elif action.nargs == ONE_OR_MORE:
+                return '%s [...]' % metavar_formatter(1)
+            elif isinstance(action.nargs, int) and action.nargs > 1:
+                return '{}{{{}}}'.format('%s' % metavar_formatter(1), action.nargs)
+
+        return super()._format_args(action, default_metavar)
 
 
 # noinspection PyCompatibility
diff --git a/tests/test_argparse_completer.py b/tests/test_argparse_completer.py
index 4313647b..151923ea 100644
--- a/tests/test_argparse_completer.py
+++ b/tests/test_argparse_completer.py
@@ -107,6 +107,10 @@ class AutoCompleteTester(cmd2.Cmd):
     ############################################################################################################
     # Begin code related to testing choices, choices_function, and choices_method parameters
     ############################################################################################################
+    STR_METAVAR = "HEADLESS"
+    TUPLE_METAVAR = ('arg1', 'others')
+    CUSTOM_DESC_HEADER = "Custom Header"
+
     def choices_method(self) -> List[str]:
         """Method that provides choices"""
         return choices_from_method
@@ -128,8 +132,14 @@ class AutoCompleteTester(cmd2.Cmd):
                                 choices_function=choices_function)
     choices_parser.add_argument("-m", "--method", help="a flag populated with a choices method",
                                 choices_method=choices_method)
-    choices_parser.add_argument('-n', "--no_header", help='this arg has a no descriptive header',
-                                choices_method=completion_item_method)
+    choices_parser.add_argument('-d', "--desc_header", help='this arg has a descriptive header',
+                                choices_method=completion_item_method,
+                                descriptive_header=CUSTOM_DESC_HEADER)
+    choices_parser.add_argument('-n', "--no_header", help='this arg has no descriptive header',
+                                choices_method=completion_item_method, metavar=STR_METAVAR)
+    choices_parser.add_argument('-t', "--tuple_metavar", help='this arg has tuple for a metavar',
+                                choices_method=completion_item_method, metavar=TUPLE_METAVAR,
+                                nargs=argparse.ONE_OR_MORE)
     choices_parser.add_argument('-i', '--int', type=int, help='a flag with an int type',
                                 choices=static_int_choices_list)
 
@@ -683,15 +693,73 @@ def test_unfinished_flag_error(ac_app, command_and_args, text, is_error, capsys)
     assert is_error == all(x in out for x in ["Error: argument", "expected"])
 
 
-def test_completion_items_default_header(ac_app):
+def test_completion_items_arg_header(ac_app):
+    # Test when metavar is None
+    text = ''
+    line = 'choices --desc_header {}'.format(text)
+    endidx = len(line)
+    begidx = endidx - len(text)
+
+    complete_tester(text, line, begidx, endidx, ac_app)
+    assert "DESC_HEADER" in ac_app.completion_header
+
+    # Test when metavar is a string
+    text = ''
+    line = 'choices --no_header {}'.format(text)
+    endidx = len(line)
+    begidx = endidx - len(text)
+
+    complete_tester(text, line, begidx, endidx, ac_app)
+    assert ac_app.STR_METAVAR in ac_app.completion_header
+
+    # Test when metavar is a tuple
+    text = ''
+    line = 'choices --tuple_metavar {}'.format(text)
+    endidx = len(line)
+    begidx = endidx - len(text)
+
+    # We are completing the first argument of this flag. The first element in the tuple should be the column header.
+    complete_tester(text, line, begidx, endidx, ac_app)
+    assert ac_app.TUPLE_METAVAR[0].upper() in ac_app.completion_header
+
+    text = ''
+    line = 'choices --tuple_metavar token_1 {}'.format(text)
+    endidx = len(line)
+    begidx = endidx - len(text)
+
+    # We are completing the second argument of this flag. The second element in the tuple should be the column header.
+    complete_tester(text, line, begidx, endidx, ac_app)
+    assert ac_app.TUPLE_METAVAR[1].upper() in ac_app.completion_header
+
+    text = ''
+    line = 'choices --tuple_metavar token_1 token_2 {}'.format(text)
+    endidx = len(line)
+    begidx = endidx - len(text)
+
+    # We are completing the third argument of this flag. It should still be the second tuple element
+    # in the column header since the tuple only has two strings in it.
+    complete_tester(text, line, begidx, endidx, ac_app)
+    assert ac_app.TUPLE_METAVAR[1].upper() in ac_app.completion_header
+
+
+def test_completion_items_descriptive_header(ac_app):
     from cmd2.argparse_completer import DEFAULT_DESCRIPTIVE_HEADER
 
+    # This argument provided a descriptive header
+    text = ''
+    line = 'choices --desc_header {}'.format(text)
+    endidx = len(line)
+    begidx = endidx - len(text)
+
+    complete_tester(text, line, begidx, endidx, ac_app)
+    assert ac_app.CUSTOM_DESC_HEADER in ac_app.completion_header
+
+    # This argument did not provide a descriptive header, so it should be DEFAULT_DESCRIPTIVE_HEADER
     text = ''
-    line = 'choices -n {}'.format(text)
+    line = 'choices --no_header {}'.format(text)
     endidx = len(line)
     begidx = endidx - len(text)
 
-    # This positional argument did not provide a descriptive header, so it should be DEFAULT_DESCRIPTIVE_HEADER
     complete_tester(text, line, begidx, endidx, ac_app)
     assert DEFAULT_DESCRIPTIVE_HEADER in ac_app.completion_header
 
diff --git a/tests/test_argparse_custom.py b/tests/test_argparse_custom.py
index f4db12b6..3ce90118 100644
--- a/tests/test_argparse_custom.py
+++ b/tests/test_argparse_custom.py
@@ -260,3 +260,10 @@ def test_override_parser():
     # Verify DEFAULT_ARGUMENT_PARSER is now our CustomParser
     from examples.custom_parser import CustomParser
     assert DEFAULT_ARGUMENT_PARSER == CustomParser
+
+
+def test_apcustom_metavar_tuple():
+    # Test the case when a tuple metavar is used with nargs an integer > 1
+    parser = Cmd2ArgumentParser()
+    parser.add_argument('--aflag', nargs=2, metavar=('foo', 'bar'), help='This is a test')
+    assert '[--aflag foo bar]' in parser.format_help()