Colorize the CLI help text

mypy/dmypy/client.py

@@ -19,20 +19,28 @@
 from mypy.dmypy_os import alive, kill
 from mypy.dmypy_util import DEFAULT_STATUS_FILE, receive
 from mypy.ipc import IPCClient, IPCException
-from mypy.util import check_python_version, get_terminal_width, should_force_color
+from mypy.util import (
+    ColoredHelpFormatter,
+    check_python_version,
+    get_terminal_width,
+    should_force_color,
+)
 from mypy.version import __version__
 
 # Argument parser.  Subparsers are tied to action functions by the
 # @action(subparse) decorator.
 
 
-class AugmentedHelpFormatter(argparse.RawDescriptionHelpFormatter):
+class AugmentedHelpFormatter(argparse.RawDescriptionHelpFormatter, ColoredHelpFormatter):
     def __init__(self, prog: str) -> None:
         super().__init__(prog=prog, max_help_position=30)
 
 
 parser = argparse.ArgumentParser(
-    prog="dmypy", description="Client for mypy daemon mode", fromfile_prefix_chars="@"
+    prog="dmypy",
+    description="Client for mypy daemon mode",
+    fromfile_prefix_chars="@",
+    formatter_class=ColoredHelpFormatter,
 )
 parser.set_defaults(action=None)
 parser.add_argument(
@@ -47,7 +55,9 @@ def __init__(self, prog: str) -> None:
 )
 subparsers = parser.add_subparsers()
 
-start_parser = p = subparsers.add_parser("start", help="Start daemon")
+start_parser = p = subparsers.add_parser(
+    "start", formatter_class=parser.formatter_class, help="Start daemon"
+)
 p.add_argument("--log-file", metavar="FILE", type=str, help="Direct daemon stdout/stderr to FILE")
 p.add_argument(
     "--timeout", metavar="TIMEOUT", type=int, help="Server shutdown timeout (in seconds)"
@@ -57,7 +67,9 @@ def __init__(self, prog: str) -> None:
 )
 
 restart_parser = p = subparsers.add_parser(
-    "restart", help="Restart daemon (stop or kill followed by start)"
+    "restart",
+    formatter_class=parser.formatter_class,
+    help="Restart daemon (stop or kill followed by start)",
 )
 p.add_argument("--log-file", metavar="FILE", type=str, help="Direct daemon stdout/stderr to FILE")
 p.add_argument(
@@ -67,13 +79,21 @@ def __init__(self, prog: str) -> None:
     "flags", metavar="FLAG", nargs="*", type=str, help="Regular mypy flags (precede with --)"
 )
 
-status_parser = p = subparsers.add_parser("status", help="Show daemon status")
+status_parser = p = subparsers.add_parser(
+    "status", formatter_class=parser.formatter_class, help="Show daemon status"
+)
 p.add_argument("-v", "--verbose", action="store_true", help="Print detailed status")
 p.add_argument("--fswatcher-dump-file", help="Collect information about the current file state")
 
-stop_parser = p = subparsers.add_parser("stop", help="Stop daemon (asks it politely to go away)")
+stop_parser = p = subparsers.add_parser(
+    "stop",
+    formatter_class=parser.formatter_class,
+    help="Stop daemon (asks it politely to go away)",
+)
 
-kill_parser = p = subparsers.add_parser("kill", help="Kill daemon (kills the process)")
+kill_parser = p = subparsers.add_parser(
+    "kill", formatter_class=parser.formatter_class, help="Kill daemon (kills the process)"
+)
 
 check_parser = p = subparsers.add_parser(
     "check", formatter_class=AugmentedHelpFormatter, help="Check some files (requires daemon)"
@@ -137,7 +157,9 @@ def __init__(self, prog: str) -> None:
 p.add_argument("--remove", metavar="FILE", nargs="*", help="Files to remove from the run")
 
 suggest_parser = p = subparsers.add_parser(
-    "suggest", help="Suggest a signature or show call sites for a specific function"
+    "suggest",
+    formatter_class=parser.formatter_class,
+    help="Suggest a signature or show call sites for a specific function",
 )
 p.add_argument(
     "function",
@@ -177,7 +199,9 @@ def __init__(self, prog: str) -> None:
 )
 
 inspect_parser = p = subparsers.add_parser(
-    "inspect", help="Locate and statically inspect expression(s)"
+    "inspect",
+    formatter_class=parser.formatter_class,
+    help="Locate and statically inspect expression(s)",
 )
 p.add_argument(
     "location",
@@ -238,9 +262,13 @@ def __init__(self, prog: str) -> None:
     help="Re-parse and re-type-check file before inspection (may be slow)",
 )
 
-hang_parser = p = subparsers.add_parser("hang", help="Hang for 100 seconds")
+hang_parser = p = subparsers.add_parser(
+    "hang", formatter_class=parser.formatter_class, help="Hang for 100 seconds"
+)
 
-daemon_parser = p = subparsers.add_parser("daemon", help="Run daemon in foreground")
+daemon_parser = p = subparsers.add_parser(
+    "daemon", formatter_class=parser.formatter_class, help="Run daemon in foreground"
+)
 p.add_argument(
     "--timeout", metavar="TIMEOUT", type=int, help="Server shutdown timeout (in seconds)"
 )
@@ -249,7 +277,7 @@ def __init__(self, prog: str) -> None:
     "flags", metavar="FLAG", nargs="*", type=str, help="Regular mypy flags (precede with --)"
 )
 p.add_argument("--options-data", help=argparse.SUPPRESS)
-help_parser = p = subparsers.add_parser("help")
+help_parser = p = subparsers.add_parser("help", formatter_class=parser.formatter_class)
 
 del p
 

mypy/main.py

@@ -67,11 +67,14 @@ def main(
         args = sys.argv[1:]
 
     fscache = FileSystemCache()
-    sources, options = process_options(args, stdout=stdout, stderr=stderr, fscache=fscache)
+    formatter = util.FancyFormatter(stdout, stderr, hide_error_codes=False)
+    sources, options = process_options(
+        args, stdout=stdout, stderr=stderr, fscache=fscache, formatter=formatter
+    )
     if clean_exit:
         options.fast_exit = False
 
-    formatter = util.FancyFormatter(stdout, stderr, options.hide_error_codes)
+    formatter.hide_error_codes = options.hide_error_codes
 
     if options.install_types and (stdout is not sys.stdout or stderr is not sys.stderr):
         # Since --install-types performs user input, we want regular stdout and stderr.
@@ -215,9 +218,9 @@ def show_messages(
 
 
 # Make the help output a little less jarring.
-class AugmentedHelpFormatter(argparse.RawDescriptionHelpFormatter):
-    def __init__(self, prog: str) -> None:
-        super().__init__(prog=prog, max_help_position=28)
+class AugmentedHelpFormatter(argparse.RawDescriptionHelpFormatter, util.ColoredHelpFormatter):
+    def __init__(self, prog: str, formatter: util.FancyFormatter) -> None:
+        super().__init__(prog=prog, max_help_position=28, formatter=formatter)
 
     def _fill_text(self, text: str, width: int, indent: str) -> str:
         if "\n" in text:
@@ -440,6 +443,7 @@ def process_options(
     fscache: FileSystemCache | None = None,
     program: str = "mypy",
     header: str = HEADER,
+    formatter: util.FancyFormatter | None = None,
 ) -> tuple[list[BuildSource], Options]:
     """Parse command line arguments.
 
@@ -448,14 +452,15 @@ def process_options(
     """
     stdout = stdout or sys.stdout
     stderr = stderr or sys.stderr
+    formatter = formatter or util.FancyFormatter(stdout, stderr, False)
 
     parser = CapturableArgumentParser(
         prog=program,
         usage=header,
         description=DESCRIPTION,
         epilog=FOOTER,
         fromfile_prefix_chars="@",
-        formatter_class=AugmentedHelpFormatter,
+        formatter_class=lambda prog: AugmentedHelpFormatter(prog, formatter),
         add_help=False,
         stdout=stdout,
         stderr=stderr,

mypy/stubgen.py

@@ -149,6 +149,7 @@
     UnionType,
     get_proper_type,
 )
+from mypy.util import ColoredHelpFormatter
 from mypy.visitor import NodeVisitor
 
 TYPING_MODULE_NAMES: Final = ("typing", "typing_extensions")
@@ -1942,7 +1943,9 @@ def generate_stubs(options: Options) -> None:
 
 
 def parse_options(args: list[str]) -> Options:
-    parser = argparse.ArgumentParser(prog="stubgen", usage=HEADER, description=DESCRIPTION)
+    parser = argparse.ArgumentParser(
+        prog="stubgen", usage=HEADER, description=DESCRIPTION, formatter_class=ColoredHelpFormatter
+    )
 
     parser.add_argument(
         "--ignore-errors",

mypy/stubtest.py

@@ -40,7 +40,13 @@
 from mypy.config_parser import parse_config_file
 from mypy.evalexpr import UNKNOWN, evaluate_expression
 from mypy.options import Options
-from mypy.util import FancyFormatter, bytes_to_human_readable_repr, is_dunder, plural_s
+from mypy.util import (
+    ColoredHelpFormatter,
+    FancyFormatter,
+    bytes_to_human_readable_repr,
+    is_dunder,
+    plural_s,
+)
 
 
 class Missing:
@@ -1906,7 +1912,8 @@ def set_strict_flags() -> None:  # not needed yet
 
 def parse_options(args: list[str]) -> _Arguments:
     parser = argparse.ArgumentParser(
-        description="Compares stubs to objects introspected from the runtime."
+        description="Compares stubs to objects introspected from the runtime.",
+        formatter_class=lambda prog: ColoredHelpFormatter(prog, formatter=_formatter),
     )
     parser.add_argument("modules", nargs="*", help="Modules to test")
     parser.add_argument(

mypy/util.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import argparse
 import hashlib
 import io
 import os
@@ -12,7 +13,7 @@
 import time
 from importlib import resources as importlib_resources
 from typing import IO, Callable, Container, Final, Iterable, Sequence, Sized, TypeVar
-from typing_extensions import Literal
+from typing_extensions import Literal, TypeAlias
 
 try:
     import curses
@@ -35,6 +36,7 @@
         TYPESHED_DIR = str(_resource.parent / "typeshed")
 
 
+Color: TypeAlias = Literal["red", "green", "blue", "yellow", "none"]
 ENCODING_RE: Final = re.compile(rb"([ \t\v]*#.*(\r\n?|\n))??[ \t\v]*#.*coding[:=][ \t]*([-\w.]+)")
 
 DEFAULT_SOURCE_OFFSET: Final = 4
@@ -630,7 +632,7 @@ def initialize_unix_colors(self) -> bool:
     def style(
         self,
         text: str,
-        color: Literal["red", "green", "blue", "yellow", "none"],
+        color: Color,
         bold: bool = False,
         underline: bool = False,
         dim: bool = False,
@@ -826,3 +828,53 @@ def quote_docstring(docstr: str) -> str:
         return f"''{docstr_repr}''"
     else:
         return f'""{docstr_repr}""'
+
+
+class ColoredHelpFormatter(argparse.HelpFormatter):
+    styles: dict[str, Color] = {"groups": "yellow", "args": "green", "metavar": "blue"}
+
+    def __init__(
+        self,
+        prog: str,
+        indent_increment: int = 2,
+        max_help_position: int = 24,
+        width: int | None = None,
+        formatter: FancyFormatter | None = None,
+    ) -> None:
+        super().__init__(prog, indent_increment, max_help_position, width)
+        self.fancy_fmt = formatter or FancyFormatter(sys.stdout, sys.stdin, False)
+
+    def start_section(self, heading: str | None) -> None:
+        if heading in {"positional arguments", "optional arguments", "options"}:
+            # make argparse generated headings consistent with mypy headings
+            heading = heading.capitalize()
+        if heading:
+            heading = self.fancy_fmt.style(heading, self.styles["groups"], bold=True)
+        return super().start_section(heading)
+
+    def _format_action(self, action: argparse.Action) -> str:
+        action_help = super()._format_action(action)
+        action_header = self._format_action_invocation(action)
+        padding, header, help = action_help.partition(action_header)
+
+        # highlight the action header
+        if not action.option_strings:  # positional-argument
+            header = self.fancy_fmt.style(header, self.styles["args"], bold=True)
+        else:  # --optional-argument METAVAR
+            parts: list[str] = []
+            for part in header.split(", "):
+                opt_str, space, metavar = part.partition(" ")
+                opt_str = self.fancy_fmt.style(opt_str, self.styles["args"], bold=True)
+                if metavar:
+                    metavar = self.fancy_fmt.style(metavar, self.styles["metavar"], bold=True)
+                parts.append(opt_str + space + metavar)
+            header = ", ".join(parts)
+
+        # highlight --optional-argument in help (may span multiple lines)
+        styled_repl = self.fancy_fmt.style(r"\1", self.styles["args"])
+        help = re.sub(
+            r"(?P<sep>^|\s)(?P<opt>--\w+(?:-(?:\n +)?\w+)*)",
+            repl=lambda m: (m.group("sep") + re.sub(r"(\S+)", styled_repl, m.group("opt"))),
+            string=help,
+        )
+        return padding + header + help