Removed custom formatting for argparse.ZERO_OR_MORE and argparse.ONE_OR_MORE.

Added rich-argparse support for coloring cmd2's custom nargs formatting.

Author: kmvanbrunt

cmd2/argparse_custom.py

@@ -258,14 +258,11 @@ def get_items(self) -> list[CompletionItems]:
 import argparse
 import re
 import sys
-from argparse import (
-    ONE_OR_MORE,
-    ZERO_OR_MORE,
-    ArgumentError,
-)
+from argparse import ArgumentError
 from collections.abc import (
     Callable,
     Iterable,
+    Iterator,
     Sequence,
 )
 from gettext import gettext
@@ -1296,29 +1293,68 @@ def format_tuple(tuple_size: int) -> tuple[str, ...]:
 
         return format_tuple
 
+    def _build_nargs_range_str(self, nargs_range: tuple[int, int | float]) -> str:
+        """Generate nargs range string for help text."""
+        if nargs_range[1] == constants.INFINITY:
+            # {min+}
+            range_str = f"{{{nargs_range[0]}+}}"
+        else:
+            # {min..max}
+            range_str = f"{{{nargs_range[0]}..{nargs_range[1]}}}"
+
+        return range_str
+
     def _format_args(self, action: argparse.Action, default_metavar: str) -> str:
-        """Handle ranged nargs and make other output less verbose."""
+        """Override to handle cmd2's custom nargs formatting.
+
+        All formats in this function need to be handled by _rich_metavar_parts().
+        """
         metavar = self._determine_metavar(action, default_metavar)
         metavar_formatter = self._metavar_formatter(action, default_metavar)
 
         # Handle nargs specified as a range
         nargs_range = action.get_nargs_range()  # type: ignore[attr-defined]
         if nargs_range is not None:
-            range_str = f'{nargs_range[0]}+' if nargs_range[1] == constants.INFINITY else f'{nargs_range[0]}..{nargs_range[1]}'
-
-            return '{}{{{}}}'.format('%s' % metavar_formatter(1), range_str)  # noqa: UP031
-
-        # 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.
-        if isinstance(metavar, str):
-            if action.nargs == ZERO_OR_MORE:
-                return '[%s [...]]' % metavar_formatter(1)  # noqa: UP031
-            if action.nargs == ONE_OR_MORE:
-                return '%s [...]' % metavar_formatter(1)  # noqa: UP031
-            if isinstance(action.nargs, int) and action.nargs > 1:
-                return '{}{{{}}}'.format('%s' % metavar_formatter(1), action.nargs)  # noqa: UP031
+            arg_str = '%s' % metavar_formatter(1)  # noqa: UP031
+            range_str = self._build_nargs_range_str(nargs_range)
+            return f"{arg_str}{range_str}"
+
+        # When nargs is just a number, argparse repeats the arg in the help text.
+        # For instance, when nargs=5 the help text looks like: 'command arg arg arg arg arg'.
+        # To make this less verbose, format it like: 'command arg{5}'.
+        # Do not customize the output when metavar is a tuple of strings. Allow argparse's
+        # formatter to handle that instead.
+        if isinstance(metavar, str) and isinstance(action.nargs, int) and action.nargs > 1:
+            arg_str = '%s' % metavar_formatter(1)  # noqa: UP031
+            return f"{arg_str}{{{action.nargs}}}"
+
+        # Fallback to parent for all other cases
+        return super()._format_args(action, default_metavar)
+
+    def _rich_metavar_parts(
+        self,
+        action: argparse.Action,
+        default_metavar: str,
+    ) -> Iterator[tuple[str, bool]]:
+        """Override to handle all cmd2-specific formatting in _format_args()."""
+        metavar = self._determine_metavar(action, default_metavar)
+        metavar_formatter = self._metavar_formatter(action, default_metavar)
 
-        return super()._format_args(action, default_metavar)  # type: ignore[arg-type]
+        # Handle nargs specified as a range
+        nargs_range = action.get_nargs_range()  # type: ignore[attr-defined]
+        if nargs_range is not None:
+            yield "%s" % metavar_formatter(1), True  # noqa: UP031
+            yield self._build_nargs_range_str(nargs_range), False
+            return
+
+        # Handle specific integer nargs (e.g., nargs=5 -> arg{5})
+        if isinstance(metavar, str) and isinstance(action.nargs, int) and action.nargs > 1:
+            yield "%s" % metavar_formatter(1), True  # noqa: UP031
+            yield f"{{{action.nargs}}}", False
+            return
+
+        # Fallback to parent for all other cases
+        yield from super()._rich_metavar_parts(action, default_metavar)
 
 
 class RawDescriptionCmd2HelpFormatter(

tests/test_argparse_custom.py

@@ -9,13 +9,9 @@
     Cmd2ArgumentParser,
     constants,
 )
-from cmd2.argparse_custom import (
-    generate_range_error,
-)
+from cmd2.argparse_custom import generate_range_error
 
-from .conftest import (
-    run_cmd,
-)
+from .conftest import run_cmd
 
 
 class ApCustomTestApp(cmd2.Cmd):
@@ -29,8 +25,6 @@ def __init__(self, *args, **kwargs) -> None:
     range_parser.add_argument('--arg1', nargs=2)
     range_parser.add_argument('--arg2', nargs=(3,))
     range_parser.add_argument('--arg3', nargs=(2, 3))
-    range_parser.add_argument('--arg4', nargs=argparse.ZERO_OR_MORE)
-    range_parser.add_argument('--arg5', nargs=argparse.ONE_OR_MORE)
 
     @cmd2.with_argparser(range_parser)
     def do_range(self, _) -> None:
@@ -89,7 +83,7 @@ def test_apcustom_usage() -> None:
 def test_apcustom_nargs_help_format(cust_app) -> None:
     out, _err = run_cmd(cust_app, 'help range')
     assert 'Usage: range [-h] [--arg0 ARG0] [--arg1 ARG1{2}] [--arg2 ARG2{3+}]' in out[0]
-    assert '             [--arg3 ARG3{2..3}] [--arg4 [ARG4 [...]]] [--arg5 ARG5 [...]]' in out[1]
+    assert '             [--arg3 ARG3{2..3}]' in out[1]
 
 
 def test_apcustom_nargs_range_validation(cust_app) -> None:
@@ -114,6 +108,50 @@ def test_apcustom_nargs_range_validation(cust_app) -> None:
     assert not err
 
 
+@pytest.mark.parametrize(
+    ('nargs', 'expected_parts'),
+    [
+        # arg{2}
+        (
+            2,
+            [("arg", True), ("{2}", False)],
+        ),
+        # arg{2+}
+        (
+            (2,),
+            [("arg", True), ("{2+}", False)],
+        ),
+        # arg{0..5}
+        (
+            (0, 5),
+            [("arg", True), ("{0..5}", False)],
+        ),
+    ],
+)
+def test_rich_metavar_parts(
+    nargs: int | tuple[int, int | float],
+    expected_parts: list[tuple[str, bool]],
+) -> None:
+    """
+    Test cmd2's override of _rich_metavar_parts which handles custom nargs formats.
+
+    :param nargs: the arguments nargs value
+    :param expected_parts: list to compare to _rich_metavar_parts's return value
+
+                           Each element in this list is a 2-item tuple.
+                           item 1: one part of the argument string outputted by _format_args
+                           item 2: boolean stating whether rich-argparse should color this part
+    """
+    parser = Cmd2ArgumentParser()
+    help_formatter = parser._get_formatter()
+
+    action = parser.add_argument("arg", nargs=nargs)  # type: ignore[arg-type]
+    default_metavar = help_formatter._get_default_metavar_for_positional(action)
+
+    parts = help_formatter._rich_metavar_parts(action, default_metavar)
+    assert list(parts) == expected_parts
+
+
 @pytest.mark.parametrize(
     'nargs_tuple',
     [
@@ -149,7 +187,7 @@ def test_apcustom_narg_tuple_zero_base() -> None:
     arg = parser.add_argument('arg', nargs=(0,))
     assert arg.nargs == argparse.ZERO_OR_MORE
     assert arg.nargs_range is None
-    assert "[arg [...]]" in parser.format_help()
+    assert "[arg ...]" in parser.format_help()
 
     parser = Cmd2ArgumentParser()
     arg = parser.add_argument('arg', nargs=(0, 1))
@@ -169,7 +207,7 @@ def test_apcustom_narg_tuple_one_base() -> None:
     arg = parser.add_argument('arg', nargs=(1,))
     assert arg.nargs == argparse.ONE_OR_MORE
     assert arg.nargs_range is None
-    assert "arg [...]" in parser.format_help()
+    assert "arg [arg ...]" in parser.format_help()
 
     parser = Cmd2ArgumentParser()
     arg = parser.add_argument('arg', nargs=(1, 5))