Improved way to create metavars coloured strings

CHANGELOG.md

@@ -2,6 +2,13 @@
 
 ## Unreleased
 
+### Fixes
+- [GH-125](https://github.com/hamdanal/rich-argparse/issues/125),
+  [GH-127](https://github.com/hamdanal/rich-argparse/pull/127),
+  [PR-128](https://github.com/hamdanal/rich-argparse/pull/128)
+  Redesign metavar styling to fix broken colors of usage when some metavars are wrapped to multiple
+  lines. The brackets and spaces of metavars are no longer colored.
+
 ## 1.5.2 - 2024-06-15
 
 ### Fixes

rich_argparse/__init__.py

@@ -305,36 +305,94 @@ def find_span(_string: str) -> tuple[int, int]:
                 usage = action.option_strings[0]
             start, end = find_span(usage)
             yield r.Span(start, end, "argparse.args")
+            pos = end + 1
             if action.nargs != 0:
-                metavar = self._format_args(action, self._get_default_metavar_for_optional(action))
-                start, end = find_span(metavar)
-                yield r.Span(start, end, "argparse.metavar")
+                default_metavar = self._get_default_metavar_for_optional(action)
+                for metavar_part, colorize in self._rich_metavar_parts(action, default_metavar):
+                    start, end = find_span(metavar_part)
+                    if colorize:
+                        yield r.Span(start, end, "argparse.metavar")
+                    pos = end
             pos = end + 1
         for action in positionals:  # positionals come at the end
-            metavar = self._get_default_metavar_for_positional(action)
-            metavar_tuple = self._metavar_formatter(action, metavar)(1)
-            usage = metavar_tuple[0]
-            if isinstance(action.nargs, int):
-                nargs = action.nargs
-            elif action.nargs in (argparse.REMAINDER, argparse.SUPPRESS):
-                nargs = 0
-            elif action.nargs in (None, argparse.OPTIONAL, argparse.PARSER):
-                nargs = 1
-            elif action.nargs == argparse.ZERO_OR_MORE:
-                if sys.version_info >= (3, 9):  # pragma: >=3.9 cover
-                    nargs = 2 if len(metavar_tuple) == 2 else 1
-                else:  # pragma: <3.9 cover
-                    nargs = 2
-            elif action.nargs == argparse.ONE_OR_MORE:
-                nargs = 2
-            else:  # pragma: no cover
-                # unknown nargs, fallback to coloring the whole thing
-                usage = self._format_args(action, metavar)
-                nargs = 1
-            for _ in range(nargs):
-                start, end = find_span(usage)
-                yield r.Span(start, end, "argparse.args")
-                pos = end + 1
+            default_metavar = self._get_default_metavar_for_positional(action)
+            for metavar_part, colorize in self._rich_metavar_parts(action, default_metavar):
+                start, end = find_span(metavar_part)
+                if colorize:
+                    yield r.Span(start, end, "argparse.args")
+                pos = end
+            pos = end + 1
+
+    def _rich_metavar_parts(
+        self, action: Action, default_metavar: str
+    ) -> Iterator[tuple[str, bool]]:
+        get_metavar = self._metavar_formatter(action, default_metavar)
+        # similar to self._format_args but yields (part, colorize) of the metavar
+        if action.nargs is None:
+            # '%s' % get_metavar(1)
+            yield "%s" % get_metavar(1), True  # noqa: UP031
+        elif action.nargs == argparse.OPTIONAL:
+            # '[%s]' % get_metavar(1)
+            yield from (
+                ("[", False),
+                ("%s" % get_metavar(1), True),  # noqa: UP031
+                ("]", False),
+            )
+        elif action.nargs == argparse.ZERO_OR_MORE:
+            if sys.version_info < (3, 9) or len(get_metavar(1)) == 2:  # pragma: <3.9 cover
+                metavar = get_metavar(2)
+                # '[%s [%s ...]]' % metavar
+                yield from (
+                    ("[", False),
+                    ("%s" % metavar[0], True),  # noqa: UP031
+                    (" [", False),
+                    ("%s" % metavar[1], True),  # noqa: UP031
+                    (" ", False),
+                    ("...", True),
+                    ("]]", False),
+                )
+            else:  # pragma: >=3.9 cover
+                # '[%s ...]' % metavar
+                yield from (
+                    ("[", False),
+                    ("%s" % get_metavar(1), True),  # noqa: UP031
+                    (" ", False),
+                    ("...", True),
+                    ("]", False),
+                )
+        elif action.nargs == argparse.ONE_OR_MORE:
+            # '%s [%s ...]' % get_metavar(2)
+            metavar = get_metavar(2)
+            yield from (
+                ("%s" % metavar[0], True),  # noqa: UP031
+                (" [", False),
+                ("%s" % metavar[1], True),  # noqa: UP031
+                (" ", False),
+                ("...", True),
+                ("]", False),
+            )
+        elif action.nargs == argparse.REMAINDER:
+            # '...'
+            yield "...", True
+        elif action.nargs == argparse.PARSER:
+            # '%s ...' % get_metavar(1)
+            yield from (
+                ("%s" % get_metavar(1), True),  # noqa: UP031
+                (" ", False),
+                ("...", True),
+            )
+        elif action.nargs == argparse.SUPPRESS:
+            # ''
+            yield "", False
+        else:
+            metavar = get_metavar(action.nargs)  # type: ignore[arg-type]
+            first = True
+            for met in metavar:
+                if first:
+                    first = False
+                else:
+                    yield " ", False
+                yield "%s" % met, True  # noqa: UP031
 
     def _rich_whitespace_sub(self, text: r.Text) -> r.Text:
         # do this `self._whitespace_matcher.sub(' ', text).strip()` but text is Text
@@ -433,8 +491,10 @@ def _rich_format_action_invocation(self, action: Action) -> r.Text:
             )
             if action.nargs != 0:
                 default = self._get_default_metavar_for_optional(action)
-                args_string = self._format_args(action, default)
-                action_header.append(" ").append(args_string, style="argparse.metavar")
+                action_header.append(" ")
+                for metavar_part, colorize in self._rich_metavar_parts(action, default):
+                    style = "argparse.metavar" if colorize else None
+                    action_header.append(metavar_part, style=style)
             return action_header
 
     def _rich_split_lines(self, text: r.Text, width: int) -> r.Lines:

tests/test_argparse.py

@@ -408,16 +408,16 @@ def test_actions_spans_in_usage():
 
     # https://github.com/python/cpython/issues/82619
     if sys.version_info < (3, 9):  # pragma: <3.9 cover
-        zom_metavar = "[\x1b[36mzom\x1b[0m [\x1b[36mzom\x1b[0m ...]]"
+        zom_metavar = "[\x1b[36mzom\x1b[0m [\x1b[36mzom\x1b[0m \x1b[36m...\x1b[0m]]"
     else:  # pragma: >=3.9 cover
-        zom_metavar = "[\x1b[36mzom\x1b[0m ...]"
+        zom_metavar = "[\x1b[36mzom\x1b[0m \x1b[36m...\x1b[0m]"
 
     usage_text = (
         f"\x1b[38;5;208mUsage:\x1b[0m \x1b[38;5;244mPROG\x1b[0m [\x1b[36m-h\x1b[0m] "
-        f"[\x1b[36m--opt\x1b[0m \x1b[38;5;36m[OPT]\x1b[0m | "
-        f"\x1b[36m--opts\x1b[0m \x1b[38;5;36mOPTS [OPTS ...]\x1b[0m]\n                "
+        f"[\x1b[36m--opt\x1b[0m [\x1b[38;5;36mOPT\x1b[0m] | "
+        f"\x1b[36m--opts\x1b[0m \x1b[38;5;36mOPTS\x1b[0m [\x1b[38;5;36mOPTS\x1b[0m \x1b[38;5;36m...\x1b[0m]]\n                "
         f"\x1b[36mrequired\x1b[0m \x1b[36mint\x1b[0m \x1b[36mint\x1b[0m [\x1b[36moptional\x1b[0m] "
-        f"{zom_metavar} \x1b[36moom\x1b[0m [\x1b[36moom\x1b[0m ...] ... \x1b[36mparser\x1b[0m ..."
+        f"{zom_metavar} \x1b[36moom\x1b[0m [\x1b[36moom\x1b[0m \x1b[36m...\x1b[0m] \x1b[36m...\x1b[0m \x1b[36mparser\x1b[0m \x1b[36m...\x1b[0m"
     )
     expected_help_output = f"""\
     {usage_text}
@@ -434,8 +434,8 @@ def test_actions_spans_in_usage():
 
     \x1b[38;5;208mOptional Arguments:\x1b[0m
       \x1b[36m-h\x1b[0m, \x1b[36m--help\x1b[0m            \x1b[39mshow this help message and exit\x1b[0m
-      \x1b[36m--opt\x1b[0m \x1b[38;5;36m[OPT]\x1b[0m
-      \x1b[36m--opts\x1b[0m \x1b[38;5;36mOPTS [OPTS ...]\x1b[0m
+      \x1b[36m--opt\x1b[0m [\x1b[38;5;36mOPT\x1b[0m]
+      \x1b[36m--opts\x1b[0m \x1b[38;5;36mOPTS\x1b[0m [\x1b[38;5;36mOPTS\x1b[0m \x1b[38;5;36m...\x1b[0m]
     """
     assert parser.format_help() == clean(expected_help_output)
 
@@ -732,7 +732,7 @@ def test_subparsers_usage():
     rich_child2 = rich_subparsers.add_parser("sp2")
     assert rich_parent.format_usage() == (
         "\x1b[38;5;208mUsage:\x1b[0m \x1b[38;5;244mPROG\x1b[0m [\x1b[36m-h\x1b[0m] "
-        "\x1b[36m{sp1,sp2}\x1b[0m ...\n"
+        "\x1b[36m{sp1,sp2}\x1b[0m \x1b[36m...\x1b[0m\n"
     )
     assert rich_child1.format_usage() == (
         "\x1b[38;5;208mUsage:\x1b[0m \x1b[38;5;244mPROG sp1\x1b[0m [\x1b[36m-h\x1b[0m]\n"
@@ -1057,3 +1057,62 @@ def test_arg_default_spans():
     """
     help_text = parser.format_help()
     assert help_text == clean(expected_help_text)
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 13), reason="Mut ex group usage wrapping broken in Python 3.13+"
+)  # CPython issue 121151 (https://github.com/python/cpython/issues/121151)
+@pytest.mark.usefixtures("force_color")
+def test_metavar_spans():  # pragma: <3.13 cover
+    # tests exotic metavars (tuples, wrapped, different nargs, etc.) in usage and help text
+    parser = argparse.ArgumentParser(
+        prog="PROG", formatter_class=lambda prog: RichHelpFormatter(prog, width=20)
+    )
+    meg = parser.add_mutually_exclusive_group()
+    meg.add_argument("--op1", metavar="MET", nargs="?")
+    meg.add_argument("--op2", metavar=("MET1", "MET2"), nargs="*")
+    meg.add_argument("--op3", nargs="*")
+    meg.add_argument("--op4", metavar=("MET1", "MET2"), nargs="+")
+    meg.add_argument("--op5", nargs="+")
+    meg.add_argument("--op6", nargs=3)
+    meg.add_argument("--op7", metavar=("MET1", "MET2", "MET3"), nargs=3)
+    help_text = parser.format_help()
+
+    op3_metavar = "[\x1b[38;5;36mOP3\x1b[0m \x1b[38;5;36m...\x1b[0m]"
+    if sys.version_info < (3, 9):  # pragma: <3.9 cover
+        op3_metavar = f"[\x1b[38;5;36mOP3\x1b[0m {op3_metavar}]"
+
+    expected_help_text = f"""\
+    \x1b[38;5;208mUsage:\x1b[0m \x1b[38;5;244mPROG\x1b[0m [\x1b[36m-h\x1b[0m]
+                [\x1b[36m--op1\x1b[0m [\x1b[38;5;36mMET\x1b[0m]
+                | \x1b[36m--op2\x1b[0m
+                [\x1b[38;5;36mMET1\x1b[0m [\x1b[38;5;36mMET2\x1b[0m \x1b[38;5;36m...\x1b[0m]]
+                | \x1b[36m--op3\x1b[0m
+                {op3_metavar}
+                | \x1b[36m--op4\x1b[0m
+                \x1b[38;5;36mMET1\x1b[0m
+                [\x1b[38;5;36mMET2\x1b[0m \x1b[38;5;36m...\x1b[0m]
+                | \x1b[36m--op5\x1b[0m
+                \x1b[38;5;36mOP5\x1b[0m
+                [\x1b[38;5;36mOP5\x1b[0m \x1b[38;5;36m...\x1b[0m]
+                | \x1b[36m--op6\x1b[0m
+                \x1b[38;5;36mOP6\x1b[0m \x1b[38;5;36mOP6\x1b[0m
+                \x1b[38;5;36mOP6\x1b[0m |
+                \x1b[36m--op7\x1b[0m
+                \x1b[38;5;36mMET1\x1b[0m
+                \x1b[38;5;36mMET2\x1b[0m
+                \x1b[38;5;36mMET3\x1b[0m]
+
+    \x1b[38;5;208mOptional Arguments:\x1b[0m
+      \x1b[36m-h\x1b[0m, \x1b[36m--help\x1b[0m
+        \x1b[39mshow this help\x1b[0m
+        \x1b[39mmessage and exit\x1b[0m
+      \x1b[36m--op1\x1b[0m [\x1b[38;5;36mMET\x1b[0m]
+      \x1b[36m--op2\x1b[0m [\x1b[38;5;36mMET1\x1b[0m [\x1b[38;5;36mMET2\x1b[0m \x1b[38;5;36m...\x1b[0m]]
+      \x1b[36m--op3\x1b[0m {op3_metavar}
+      \x1b[36m--op4\x1b[0m \x1b[38;5;36mMET1\x1b[0m [\x1b[38;5;36mMET2\x1b[0m \x1b[38;5;36m...\x1b[0m]
+      \x1b[36m--op5\x1b[0m \x1b[38;5;36mOP5\x1b[0m [\x1b[38;5;36mOP5\x1b[0m \x1b[38;5;36m...\x1b[0m]
+      \x1b[36m--op6\x1b[0m \x1b[38;5;36mOP6\x1b[0m \x1b[38;5;36mOP6\x1b[0m \x1b[38;5;36mOP6\x1b[0m
+      \x1b[36m--op7\x1b[0m \x1b[38;5;36mMET1\x1b[0m \x1b[38;5;36mMET2\x1b[0m \x1b[38;5;36mMET3\x1b[0m
+    """
+    assert help_text == clean(expected_help_text)