Add theming to argparse

Author: hugovk

Lib/_colorize.py

@@ -135,6 +135,22 @@ def __iter__(self) -> Iterator[str]:
         return iter(self.__dataclass_fields__)
 
 
+@dataclass(frozen=True)
+class Argparse(ThemeSection):
+    usage: str = ANSIColors.BOLD_BLUE
+    prog: str = ANSIColors.BOLD_MAGENTA
+    prog_extra: str = ANSIColors.MAGENTA
+    heading: str = ANSIColors.BOLD_BLUE
+    summary_long_option: str = ANSIColors.CYAN
+    summary_short_option: str = ANSIColors.GREEN
+    summary_label: str = ANSIColors.YELLOW
+    long_option: str = ANSIColors.BOLD_CYAN
+    short_option: str = ANSIColors.BOLD_GREEN
+    label: str = ANSIColors.BOLD_YELLOW
+    action: str = ANSIColors.BOLD_GREEN
+    reset: str = ANSIColors.RESET
+
+
 @dataclass(frozen=True)
 class REPL(ThemeSection):
     prompt: str = ANSIColors.BOLD_MAGENTA
@@ -177,13 +193,15 @@ class Theme:
     When adding a new one, remember to also modify `copy_with` and `no_colors`
     below.
     """
+    argparse: Argparse = field(default_factory=Argparse)
     repl: REPL = field(default_factory=REPL)
     traceback: Traceback = field(default_factory=Traceback)
     unittest: Unittest = field(default_factory=Unittest)
 
     def copy_with(
         self,
         *,
+        argparse: Argparse | None = None,
         repl: REPL | None = None,
         traceback: Traceback | None = None,
         unittest: Unittest | None = None,
@@ -194,6 +212,7 @@ def copy_with(
         could lead to invalid terminal states.
         """
         return type(self)(
+            argparse=argparse or self.argparse,
             repl=repl or self.repl,
             traceback=traceback or self.traceback,
             unittest=unittest or self.unittest,
@@ -207,6 +226,7 @@ def no_colors(self) -> Self:
         and possible, and empty strings otherwise.
         """
         return type(self)(
+            argparse=self.argparse.no_colors(),
             repl=self.repl.no_colors(),
             traceback=self.traceback.no_colors(),
             unittest=self.unittest.no_colors(),

Lib/argparse.py

@@ -176,13 +176,13 @@ def __init__(
             width = shutil.get_terminal_size().columns
             width -= 2
 
-        from _colorize import ANSIColors, NoColors, can_colorize, decolor
+        from _colorize import can_colorize, decolor, get_theme
 
         if color and can_colorize():
-            self._ansi = ANSIColors()
+            self._theme = get_theme(force_color=True).argparse
             self._decolor = decolor
         else:
-            self._ansi = NoColors
+            self._theme = get_theme(force_no_color=True).argparse
             self._decolor = lambda text: text
 
         self._prefix_chars = prefix_chars
@@ -237,14 +237,12 @@ def format_help(self):
 
             # add the heading if the section was non-empty
             if self.heading is not SUPPRESS and self.heading is not None:
-                bold_blue = self.formatter._ansi.BOLD_BLUE
-                reset = self.formatter._ansi.RESET
-
                 current_indent = self.formatter._current_indent
                 heading_text = _('%(heading)s:') % dict(heading=self.heading)
+                t = self.formatter._theme
                 heading = (
                     f'{" " * current_indent}'
-                    f'{bold_blue}{heading_text}{reset}\n'
+                    f'{t.heading}{heading_text}{t.reset}\n'
                 )
             else:
                 heading = ''
@@ -314,26 +312,23 @@ def _join_parts(self, part_strings):
                         if part and part is not SUPPRESS])
 
     def _format_usage(self, usage, actions, groups, prefix):
-        bold_blue = self._ansi.BOLD_BLUE
-        bold_magenta = self._ansi.BOLD_MAGENTA
-        magenta = self._ansi.MAGENTA
-        reset = self._ansi.RESET
+        t = self._theme
 
         if prefix is None:
             prefix = _('usage: ')
 
         # if usage is specified, use that
         if usage is not None:
             usage = (
-                magenta
+                t.prog_extra
                 + usage
-                % {"prog": f"{bold_magenta}{self._prog}{reset}{magenta}"}
-                + reset
+                % {"prog": f"{t.prog}{self._prog}{t.reset}{t.prog_extra}"}
+                + t.reset
             )
 
         # if no optionals or positionals are available, usage is just prog
         elif usage is None and not actions:
-            usage = f"{bold_magenta}{self._prog}{reset}"
+            usage = f"{t.prog}{self._prog}{t.reset}"
 
         # if optionals and positionals are available, calculate usage
         elif usage is None:
@@ -411,10 +406,10 @@ def get_lines(parts, indent, prefix=None):
                 usage = '\n'.join(lines)
 
             usage = usage.removeprefix(prog)
-            usage = f"{bold_magenta}{prog}{reset}{usage}"
+            usage = f"{t.prog}{prog}{t.reset}{usage}"
 
         # prefix with 'usage:'
-        return f'{bold_blue}{prefix}{reset}{usage}\n\n'
+        return f'{t.usage}{prefix}{t.reset}{usage}\n\n'
 
     def _format_actions_usage(self, actions, groups):
         return ' '.join(self._get_actions_usage_parts(actions, groups))
@@ -452,10 +447,7 @@ def _get_actions_usage_parts(self, actions, groups):
 
         # collect all actions format strings
         parts = []
-        cyan = self._ansi.CYAN
-        green = self._ansi.GREEN
-        yellow = self._ansi.YELLOW
-        reset = self._ansi.RESET
+        t = self._theme
         for action in actions:
 
             # suppressed arguments are marked with None
@@ -465,7 +457,11 @@ def _get_actions_usage_parts(self, actions, groups):
             # produce all arg strings
             elif not action.option_strings:
                 default = self._get_default_metavar_for_positional(action)
-                part = green + self._format_args(action, default) + reset
+                part = (
+                    t.summary_short_option
+                    + self._format_args(action, default)
+                    + t.reset
+                )
 
                 # if it's in a group, strip the outer []
                 if action in group_actions:
@@ -481,20 +477,23 @@ def _get_actions_usage_parts(self, actions, groups):
                 if action.nargs == 0:
                     part = action.format_usage()
                     if self._is_long_option(part):
-                        part = f"{cyan}{part}{reset}"
+                        part = f"{t.summary_long_option}{part}{t.reset}"
                     elif self._is_short_option(part):
-                        part = f"{green}{part}{reset}"
+                        part = f"{t.summary_short_option}{part}{t.reset}"
 
                 # if the Optional takes a value, format is:
                 #    -s ARGS or --long ARGS
                 else:
                     default = self._get_default_metavar_for_optional(action)
                     args_string = self._format_args(action, default)
                     if self._is_long_option(option_string):
-                        option_string = f"{cyan}{option_string}"
+                        option_color = t.summary_long_option
                     elif self._is_short_option(option_string):
-                        option_string = f"{green}{option_string}"
-                    part = f"{option_string} {yellow}{args_string}{reset}"
+                        option_color = t.summary_short_option
+                    part = (
+                        f"{option_color}{option_string} "
+                        f"{t.summary_label}{args_string}{t.reset}"
+                    )
 
                 # make it look optional if it's not required or in a group
                 if not action.required and action not in group_actions:
@@ -590,17 +589,14 @@ def _format_action(self, action):
         return self._join_parts(parts)
 
     def _format_action_invocation(self, action):
-        bold_green = self._ansi.BOLD_GREEN
-        bold_cyan = self._ansi.BOLD_CYAN
-        bold_yellow = self._ansi.BOLD_YELLOW
-        reset = self._ansi.RESET
+        t = self._theme
 
         if not action.option_strings:
             default = self._get_default_metavar_for_positional(action)
             return (
-                bold_green
+                t.action
                 + ' '.join(self._metavar_formatter(action, default)(1))
-                + reset
+                + t.reset
             )
 
         else:
@@ -609,9 +605,9 @@ def color_option_strings(strings):
                 parts = []
                 for s in strings:
                     if self._is_long_option(s):
-                        parts.append(f"{bold_cyan}{s}{reset}")
+                        parts.append(f"{t.long_option}{s}{t.reset}")
                     elif self._is_short_option(s):
-                        parts.append(f"{bold_green}{s}{reset}")
+                        parts.append(f"{t.short_option}{s}{t.reset}")
                     else:
                         parts.append(s)
                 return parts
@@ -628,7 +624,7 @@ def color_option_strings(strings):
                 default = self._get_default_metavar_for_optional(action)
                 option_strings = color_option_strings(action.option_strings)
                 args_string = (
-                    f"{bold_yellow}{self._format_args(action, default)}{reset}"
+                    f"{t.label}{self._format_args(action, default)}{t.reset}"
                 )
                 return ', '.join(option_strings) + ' ' + args_string