Added ability to specify nargs ranges with no upper bound

4 files changed, 91 insertions, 58 deletions

diff --git a/cmd2/argparse_completer.py b/cmd2/argparse_completer.py
index 875fb3db..737286c1 100644
--- a/cmd2/argparse_completer.py
+++ b/cmd2/argparse_completer.py
@@ -14,7 +14,7 @@ from . import cmd2
 from . import utils
 from .ansi import ansi_safe_wcswidth, style_error
 from .argparse_custom import ATTR_SUPPRESS_TAB_HINT, ATTR_DESCRIPTIVE_COMPLETION_HEADER, ATTR_NARGS_RANGE
-from .argparse_custom import ChoicesCallable, CompletionItem, ATTR_CHOICES_CALLABLE
+from .argparse_custom import ChoicesCallable, CompletionItem, ATTR_CHOICES_CALLABLE, INFINITY, generate_range_error
 from .rl_utils import rl_force_redisplay
 
 # If no descriptive header is supplied, then this will be used instead
@@ -83,10 +83,10 @@ class AutoCompleter(object):
                 self.max = 1
             elif self.action.nargs == argparse.ZERO_OR_MORE or self.action.nargs == argparse.REMAINDER:
                 self.min = 0
-                self.max = float('inf')
+                self.max = INFINITY
             elif self.action.nargs == argparse.ONE_OR_MORE:
                 self.min = 1
-                self.max = float('inf')
+                self.max = INFINITY
             else:
                 self.min = self.action.nargs
                 self.max = self.action.nargs
@@ -553,21 +553,7 @@ class AutoCompleter(object):
 
         out_str = "\nError:\n"
         out_str += '  {0: <{width}}    '.format(prefix, width=20)
-        out_str += "Flag requires "
-
-        # This handles ONE_OR_MORE
-        if flag_arg_state.max == float('inf'):
-            out_str += "at least {} argument".format(flag_arg_state.min)
-        else:
-            if flag_arg_state.min == flag_arg_state.max:
-                out_str += "{} ".format(flag_arg_state.min)
-            else:
-                out_str += "{} to {} ".format(flag_arg_state.min, flag_arg_state.max)
-
-            if flag_arg_state.max == 1:
-                out_str += "argument"
-            else:
-                out_str += "arguments"
+        out_str += generate_range_error(flag_arg_state.min, flag_arg_state.max)
 
         out_str += ' ({} entered)'.format(flag_arg_state.count)
         print(style_error('{}\n'.format(out_str)))
diff --git a/cmd2/argparse_custom.py b/cmd2/argparse_custom.py
index 5bcbc91a..9e6805aa 100644
--- a/cmd2/argparse_custom.py
+++ b/cmd2/argparse_custom.py
@@ -14,9 +14,14 @@ whereas any other parser class won't be as explicit in their output.
 # Added capabilities
 ############################################################################################################
 
-Extends argparse nargs functionality by allowing tuples which specify a range (min, max)
+Extends argparse nargs functionality by allowing tuples which specify a range (min, max). To specify a max
+value with no upper bound, use a 1-item tuple (min,)
+
     Example:
-        The following command says the -f argument expects between 3 and 5 values (inclusive)
+        # -f argument expects at least 3 values
+        parser.add_argument('-f', nargs=(3,))
+
+        # -f argument expects 3 to 5 values
         parser.add_argument('-f', nargs=(3, 5))
 
 Tab Completion:
@@ -153,6 +158,9 @@ from .ansi import ansi_aware_write, style_error
 # The following are names of custom argparse argument attributes added by cmd2
 ############################################################################################################
 
+# Used in nargs ranges to signify there is no maximum
+INFINITY = float('inf')
+
 # A tuple specifying nargs as a range (min, max)
 ATTR_NARGS_RANGE = 'nargs_range'
 
@@ -167,6 +175,27 @@ ATTR_SUPPRESS_TAB_HINT = 'suppress_tab_hint'
 ATTR_DESCRIPTIVE_COMPLETION_HEADER = 'desc_completion_header'
 
 
+def generate_range_error(range_min: int, range_max: Union[int, float]) -> str:
+    """Generate an error message when the the number of arguments provided is not within the expected range"""
+    err_str = "expected "
+
+    if range_max == INFINITY:
+        err_str += "at least {} argument".format(range_min)
+
+        if range_min != 1:
+            err_str += "s"
+    else:
+        if range_min == range_max:
+            err_str += "{} argument".format(range_min)
+        else:
+            err_str += "{} to {} argument".format(range_min, range_max)
+
+        if range_max != 1:
+            err_str += "s"
+
+    return err_str
+
+
 class CompletionItem(str):
     """
     Completion item with descriptive text attached
@@ -218,7 +247,7 @@ orig_actions_container_add_argument = argparse._ActionsContainer.add_argument
 
 
 def _add_argument_wrapper(self, *args,
-                          nargs: Union[int, str, Tuple[int, int], None] = None,
+                          nargs: Union[int, str, Tuple[int], Tuple[int, int], None] = None,
                           choices_function: Optional[Callable[[], Iterable[Any]]] = None,
                           choices_method: Optional[Callable[[Any], Iterable[Any]]] = None,
                           completer_function: Optional[Callable[[str, str, int, int], List[str]]] = None,
@@ -235,6 +264,7 @@ def _add_argument_wrapper(self, *args,
 
     # Customized arguments from original function
     :param nargs: extends argparse nargs functionality by allowing tuples which specify a range (min, max)
+                  to specify a max value with no upper bound, use a 1-item tuple (min,)
 
     # Added args used by AutoCompleter
     :param choices_function: function that provides choices for this argument
@@ -265,9 +295,14 @@ def _add_argument_wrapper(self, *args,
         # Check if nargs was given as a range
         if isinstance(nargs, tuple):
 
+            # Handle 1-item tuple by setting max to INFINITY
+            if len(nargs) == 1:
+                nargs = (nargs[0], INFINITY)
+
             # Validate nargs tuple
-            if len(nargs) != 2 or not isinstance(nargs[0], int) or not isinstance(nargs[1], int):
-                raise ValueError('Ranged values for nargs must be a tuple of 2 integers')
+            if len(nargs) != 2 or not isinstance(nargs[0], int) or \
+                    not (isinstance(nargs[1], int) or nargs[1] == INFINITY):
+                raise ValueError('Ranged values for nargs must be a tuple of 1 or 2 integers')
             if nargs[0] >= nargs[1]:
                 raise ValueError('Invalid nargs range. The first value must be less than the second')
             if nargs[0] < 0:
@@ -275,13 +310,26 @@ def _add_argument_wrapper(self, *args,
 
             # Save the nargs tuple as our range setting
             nargs_range = nargs
+            range_min = nargs_range[0]
+            range_max = nargs_range[1]
 
             # Convert nargs into a format argparse recognizes
-            if nargs_range[0] == 0:
-                if nargs_range[1] > 1:
-                    nargs_adjusted = argparse.ZERO_OR_MORE
-                else:
+            if range_min == 0:
+                if range_max == 1:
                     nargs_adjusted = argparse.OPTIONAL
+
+                    # No range needed since (0, 1) is just argparse.OPTIONAL
+                    nargs_range = None
+                else:
+                    nargs_adjusted = argparse.ZERO_OR_MORE
+                    if range_max == INFINITY:
+                        # No range needed since (0, INFINITY) is just argparse.ZERO_OR_MORE
+                        nargs_range = None
+            elif range_min == 1 and range_max == INFINITY:
+                nargs_adjusted = argparse.ONE_OR_MORE
+
+                # No range needed since (1, INFINITY) is just argparse.ONE_OR_MORE
+                nargs_range = None
             else:
                 nargs_adjusted = argparse.ONE_OR_MORE
         else:
@@ -342,7 +390,12 @@ def _get_nargs_pattern_wrapper(self, action) -> str:
     # Wrapper around ArgumentParser._get_nargs_pattern behavior to support nargs ranges
     nargs_range = getattr(action, ATTR_NARGS_RANGE, None)
     if nargs_range is not None:
-        nargs_pattern = '(-*A{{{},{}}}-*)'.format(nargs_range[0], nargs_range[1])
+        if nargs_range[1] == INFINITY:
+            range_max = ''
+        else:
+            range_max = nargs_range[1]
+
+        nargs_pattern = '(-*A{{{},{}}}-*)'.format(nargs_range[0], range_max)
 
         # if this is an optional action, -- is not allowed
         if action.option_strings:
@@ -375,8 +428,7 @@ def _match_argument_wrapper(self, action, arg_strings_pattern) -> int:
     if match is None:
         nargs_range = getattr(action, ATTR_NARGS_RANGE, None)
         if nargs_range is not None:
-            raise ArgumentError(action,
-                                'Expected between {} and {} arguments'.format(nargs_range[0], nargs_range[1]))
+            raise ArgumentError(action, generate_range_error(nargs_range[0], nargs_range[1]))
 
     return orig_argument_parser_match_argument(self, action, arg_strings_pattern)
 
@@ -563,7 +615,12 @@ class Cmd2HelpFormatter(argparse.RawTextHelpFormatter):
         nargs_range = getattr(action, ATTR_NARGS_RANGE, None)
 
         if nargs_range is not None:
-            result = '{}{{{}..{}}}'.format('%s' % get_metavar(1), nargs_range[0], nargs_range[1])
+            if nargs_range[1] == 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:
diff --git a/tests/test_argparse_completer.py b/tests/test_argparse_completer.py
index 1262b9e1..4ad4c560 100644
--- a/tests/test_argparse_completer.py
+++ b/tests/test_argparse_completer.py
@@ -596,10 +596,7 @@ def test_unfinished_flag_error(ac_app, command_and_args, text, is_error, capsys)
     complete_tester(text, line, begidx, endidx, ac_app)
 
     out, err = capsys.readouterr()
-    if is_error:
-        assert "Flag requires" in out
-    else:
-        assert "Flag requires" not in out
+    assert is_error == all(x in out for x in ["Error:\n", "expected"])
 
 
 def test_completion_items_default_header(ac_app):
@@ -651,11 +648,7 @@ def test_autocomp_hint(ac_app, command_and_args, text, has_hint, capsys):
 
     complete_tester(text, line, begidx, endidx, ac_app)
     out, err = capsys.readouterr()
-
-    if has_hint:
-        assert "Hint" in out
-    else:
-        assert "Hint" not in out
+    assert has_hint == ("Hint:\n" in out)
 
 
 def test_autocomp_hint_multiple_lines(ac_app, capsys):
diff --git a/tests/test_argparse_custom.py b/tests/test_argparse_custom.py
index b738efa3..17fd8334 100644
--- a/tests/test_argparse_custom.py
+++ b/tests/test_argparse_custom.py
@@ -70,28 +70,20 @@ def test_apcustom_nargs_help_format(cust_app):
 
 def test_apcustom_nargs_not_enough(cust_app):
     out, err = run_cmd(cust_app, 'range --arg1 one')
-    assert 'Error: argument --arg1: Expected between 2 and 3 arguments' in err[2]
+    assert 'Error: argument --arg1: expected 2 to 3 arguments' in err[2]
 
 
-def test_apcustom_narg_empty_tuple():
-    with pytest.raises(ValueError) as excinfo:
-        parser = cmd2.ArgParser(prog='test')
-        parser.add_argument('invalid_tuple', nargs=())
-    assert 'Ranged values for nargs must be a tuple of 2 integers' in str(excinfo.value)
-
-
-def test_apcustom_narg_single_tuple():
-    with pytest.raises(ValueError) as excinfo:
-        parser = cmd2.ArgParser(prog='test')
-        parser.add_argument('invalid_tuple', nargs=(1,))
-    assert 'Ranged values for nargs must be a tuple of 2 integers' in str(excinfo.value)
-
-
-def test_apcustom_narg_tuple_triple():
+@pytest.mark.parametrize('nargs_tuple', [
+    (),
+    ('f', 5),
+    (5, 'f'),
+    (1, 2, 3),
+])
+def test_apcustom_narg_invalid_tuples(nargs_tuple):
     with pytest.raises(ValueError) as excinfo:
         parser = cmd2.ArgParser(prog='test')
-        parser.add_argument('invalid_tuple', nargs=(1, 2, 3))
-    assert 'Ranged values for nargs must be a tuple of 2 integers' in str(excinfo.value)
+        parser.add_argument('invalid_tuple', nargs=nargs_tuple)
+    assert 'Ranged values for nargs must be a tuple of 1 or 2 integers' in str(excinfo.value)
 
 
 def test_apcustom_narg_tuple_order():
@@ -113,6 +105,11 @@ def test_apcustom_narg_tuple_zero_base():
     parser.add_argument('tuple', nargs=(0, 3))
 
 
+def test_apcustom_narg_single_tuple():
+    parser = cmd2.ArgParser(prog='test')
+    parser.add_argument('tuple', nargs=(5,))
+
+
 def test_apcustom_narg_tuple_zero_to_one():
     parser = cmd2.ArgParser(prog='test')
     parser.add_argument('tuple', nargs=(0, 1))