gh-133346: Make theming support in _colorize extensible (GH-133347)

Co-authored-by: Hugo van Kemenade <[email protected]>

Author: ambv

Doc/whatsnew/3.14.rst

@@ -1466,7 +1466,7 @@ pdb
 * Source code displayed in :mod:`pdb` will be syntax-highlighted. This feature
   can be controlled using the same methods as PyREPL, in addition to the newly
   added ``colorize`` argument of :class:`pdb.Pdb`.
-  (Contributed by Tian Gao in :gh:`133355`.)
+  (Contributed by Tian Gao and Łukasz Langa in :gh:`133355`.)
 
 
 pickle

Lib/_colorize.py

@@ -1,28 +1,17 @@
-from __future__ import annotations
 import io
 import os
 import sys
 
+from collections.abc import Callable, Iterator, Mapping
+from dataclasses import dataclass, field, Field
+
 COLORIZE = True
 
+
 # types
 if False:
-    from typing import IO, Literal
-
-    type ColorTag = Literal[
-        "PROMPT",
-        "KEYWORD",
-        "BUILTIN",
-        "COMMENT",
-        "STRING",
-        "NUMBER",
-        "OP",
-        "DEFINITION",
-        "SOFT_KEYWORD",
-        "RESET",
-    ]
-
-    theme: dict[ColorTag, str]
+    from typing import IO, Self, ClassVar
+    _theme: Theme
 
 
 class ANSIColors:
@@ -86,6 +75,186 @@ class ANSIColors:
         setattr(NoColors, attr, "")
 
 
+#
+# Experimental theming support (see gh-133346)
+#
+
+# - Create a theme by copying an existing `Theme` with one or more sections
+#   replaced, using `default_theme.copy_with()`;
+# - create a theme section by copying an existing `ThemeSection` with one or
+#   more colors replaced, using for example `default_theme.syntax.copy_with()`;
+# - create a theme from scratch by instantiating a `Theme` data class with
+#   the required sections (which are also dataclass instances).
+#
+# Then call `_colorize.set_theme(your_theme)` to set it.
+#
+# Put your theme configuration in $PYTHONSTARTUP for the interactive shell,
+# or sitecustomize.py in your virtual environment or Python installation for
+# other uses.  Your applications can call `_colorize.set_theme()` too.
+#
+# Note that thanks to the dataclasses providing default values for all fields,
+# creating a new theme or theme section from scratch is possible without
+# specifying all keys.
+#
+# For example, here's a theme that makes punctuation and operators less prominent:
+#
+#   try:
+#       from _colorize import set_theme, default_theme, Syntax, ANSIColors
+#   except ImportError:
+#       pass
+#   else:
+#       theme_with_dim_operators = default_theme.copy_with(
+#           syntax=Syntax(op=ANSIColors.INTENSE_BLACK),
+#       )
+#       set_theme(theme_with_dim_operators)
+#       del set_theme, default_theme, Syntax, ANSIColors, theme_with_dim_operators
+#
+# Guarding the import ensures that your .pythonstartup file will still work in
+# Python 3.13 and older. Deleting the variables ensures they don't remain in your
+# interactive shell's global scope.
+
+class ThemeSection(Mapping[str, str]):
+    """A mixin/base class for theme sections.
+
+    It enables dictionary access to a section, as well as implements convenience
+    methods.
+    """
+
+    # The two types below are just that: types to inform the type checker that the
+    # mixin will work in context of those fields existing
+    __dataclass_fields__: ClassVar[dict[str, Field[str]]]
+    _name_to_value: Callable[[str], str]
+
+    def __post_init__(self) -> None:
+        name_to_value = {}
+        for color_name in self.__dataclass_fields__:
+            name_to_value[color_name] = getattr(self, color_name)
+        super().__setattr__('_name_to_value', name_to_value.__getitem__)
+
+    def copy_with(self, **kwargs: str) -> Self:
+        color_state: dict[str, str] = {}
+        for color_name in self.__dataclass_fields__:
+            color_state[color_name] = getattr(self, color_name)
+        color_state.update(kwargs)
+        return type(self)(**color_state)
+
+    @classmethod
+    def no_colors(cls) -> Self:
+        color_state: dict[str, str] = {}
+        for color_name in cls.__dataclass_fields__:
+            color_state[color_name] = ""
+        return cls(**color_state)
+
+    def __getitem__(self, key: str) -> str:
+        return self._name_to_value(key)
+
+    def __len__(self) -> int:
+        return len(self.__dataclass_fields__)
+
+    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
+    summary_action: str = ANSIColors.GREEN
+    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 Syntax(ThemeSection):
+    prompt: str = ANSIColors.BOLD_MAGENTA
+    keyword: str = ANSIColors.BOLD_BLUE
+    builtin: str = ANSIColors.CYAN
+    comment: str = ANSIColors.RED
+    string: str = ANSIColors.GREEN
+    number: str = ANSIColors.YELLOW
+    op: str = ANSIColors.RESET
+    definition: str = ANSIColors.BOLD
+    soft_keyword: str = ANSIColors.BOLD_BLUE
+    reset: str = ANSIColors.RESET
+
+
+@dataclass(frozen=True)
+class Traceback(ThemeSection):
+    type: str = ANSIColors.BOLD_MAGENTA
+    message: str = ANSIColors.MAGENTA
+    filename: str = ANSIColors.MAGENTA
+    line_no: str = ANSIColors.MAGENTA
+    frame: str = ANSIColors.MAGENTA
+    error_highlight: str = ANSIColors.BOLD_RED
+    error_range: str = ANSIColors.RED
+    reset: str = ANSIColors.RESET
+
+
+@dataclass(frozen=True)
+class Unittest(ThemeSection):
+    passed: str = ANSIColors.GREEN
+    warn: str = ANSIColors.YELLOW
+    fail: str = ANSIColors.RED
+    fail_info: str = ANSIColors.BOLD_RED
+    reset: str = ANSIColors.RESET
+
+
+@dataclass(frozen=True)
+class Theme:
+    """A suite of themes for all sections of Python.
+
+    When adding a new one, remember to also modify `copy_with` and `no_colors`
+    below.
+    """
+    argparse: Argparse = field(default_factory=Argparse)
+    syntax: Syntax = field(default_factory=Syntax)
+    traceback: Traceback = field(default_factory=Traceback)
+    unittest: Unittest = field(default_factory=Unittest)
+
+    def copy_with(
+        self,
+        *,
+        argparse: Argparse | None = None,
+        syntax: Syntax | None = None,
+        traceback: Traceback | None = None,
+        unittest: Unittest | None = None,
+    ) -> Self:
+        """Return a new Theme based on this instance with some sections replaced.
+
+        Themes are immutable to protect against accidental modifications that
+        could lead to invalid terminal states.
+        """
+        return type(self)(
+            argparse=argparse or self.argparse,
+            syntax=syntax or self.syntax,
+            traceback=traceback or self.traceback,
+            unittest=unittest or self.unittest,
+        )
+
+    @classmethod
+    def no_colors(cls) -> Self:
+        """Return a new Theme where colors in all sections are empty strings.
+
+        This allows writing user code as if colors are always used. The color
+        fields will be ANSI color code strings when colorization is desired
+        and possible, and empty strings otherwise.
+        """
+        return cls(
+            argparse=Argparse.no_colors(),
+            syntax=Syntax.no_colors(),
+            traceback=Traceback.no_colors(),
+            unittest=Unittest.no_colors(),
+        )
+
+
 def get_colors(
     colorize: bool = False, *, file: IO[str] | IO[bytes] | None = None
 ) -> ANSIColors:
@@ -138,26 +307,40 @@ def can_colorize(*, file: IO[str] | IO[bytes] | None = None) -> bool:
         return hasattr(file, "isatty") and file.isatty()
 
 
-def set_theme(t: dict[ColorTag, str] | None = None) -> None:
-    global theme
+default_theme = Theme()
+theme_no_color = default_theme.no_colors()
+
+
+def get_theme(
+    *,
+    tty_file: IO[str] | IO[bytes] | None = None,
+    force_color: bool = False,
+    force_no_color: bool = False,
+) -> Theme:
+    """Returns the currently set theme, potentially in a zero-color variant.
+
+    In cases where colorizing is not possible (see `can_colorize`), the returned
+    theme contains all empty strings in all color definitions.
+    See `Theme.no_colors()` for more information.
+
+    It is recommended not to cache the result of this function for extended
+    periods of time because the user might influence theme selection by
+    the interactive shell, a debugger, or application-specific code. The
+    environment (including environment variable state and console configuration
+    on Windows) can also change in the course of the application life cycle.
+    """
+    if force_color or (not force_no_color and can_colorize(file=tty_file)):
+        return _theme
+    return theme_no_color
+
+
+def set_theme(t: Theme) -> None:
+    global _theme
 
-    if t:
-        theme = t
-        return
+    if not isinstance(t, Theme):
+        raise ValueError(f"Expected Theme object, found {t}")
 
-    colors = get_colors()
-    theme = {
-        "PROMPT": colors.BOLD_MAGENTA,
-        "KEYWORD": colors.BOLD_BLUE,
-        "BUILTIN": colors.CYAN,
-        "COMMENT": colors.RED,
-        "STRING": colors.GREEN,
-        "NUMBER": colors.YELLOW,
-        "OP": colors.RESET,
-        "DEFINITION": colors.BOLD,
-        "SOFT_KEYWORD": colors.BOLD_BLUE,
-        "RESET": colors.RESET,
-    }
+    _theme = t
 
 
-set_theme()
+set_theme(default_theme)

Lib/_pyrepl/reader.py

@@ -28,7 +28,7 @@
 from dataclasses import dataclass, field, fields
 
 from . import commands, console, input
-from .utils import wlen, unbracket, disp_str, gen_colors
+from .utils import wlen, unbracket, disp_str, gen_colors, THEME
 from .trace import trace
 
 
@@ -491,11 +491,8 @@ def get_prompt(self, lineno: int, cursor_on_line: bool) -> str:
             prompt = self.ps1
 
         if self.can_colorize:
-            prompt = (
-                f"{_colorize.theme["PROMPT"]}"
-                f"{prompt}"
-                f"{_colorize.theme["RESET"]}"
-            )
+            t = THEME()
+            prompt = f"{t.prompt}{prompt}{t.reset}"
         return prompt
 
     def push_input_trans(self, itrans: input.KeymapTranslator) -> None: