gh-133346: Make theming support in _colorize extensible

ambv · ambv · commit 33440084bf7b · 2025-05-03T17:14:13.000+02:00
diff --git a/Lib/_colorize.py b/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,68 @@ class ANSIColors:
         setattr(NoColors, attr, "")
 
 
+class ThemeSection(Mapping[str, str]):
+    __dataclass_fields__: ClassVar[dict[str, Field[str]]]
+    _name_to_value: Callable[[str], str]
+
+    def __post_init__(self) -> None:
+        name_to_value = {}
+        for color_name, color_field in self.__dataclass_fields__.items():
+            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, color_field in self.__dataclass_fields__.items():
+            color_state[color_name] = getattr(self, color_name)
+        color_state.update(kwargs)
+        return type(self)(**color_state)
+
+    def no_colors(self) -> Self:
+        color_state: dict[str, str] = {}
+        for color_name, color_field in self.__dataclass_fields__.items():
+            color_state[color_name] = ""
+        return type(self)(**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 REPL(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 Theme:
+    repl: REPL = field(default_factory=REPL)
+
+    def copy_with(self, *, repl: REPL | None) -> Self:
+        return type(self)(
+            repl=repl or self.repl,
+        )
+
+    def no_colors(self) -> Self:
+        return type(self)(
+            repl=self.repl.no_colors(),
+        )
+
+
 def get_colors(
     colorize: bool = False, *, file: IO[str] | IO[bytes] | None = None
 ) -> ANSIColors:
@@ -138,26 +189,21 @@ 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) -> Theme:
+    if can_colorize(file=tty_file):
+        return _theme
+    return theme_no_color
+
 
-    if t:
-        theme = t
-        return
+def set_theme(t: Theme) -> None:
+    global _theme
 
-    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
+    return
 
 
-set_theme()
+set_theme(default_theme)
diff --git a/Lib/_pyrepl/reader.py b/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,10 +491,11 @@ def get_prompt(self, lineno: int, cursor_on_line: bool) -> str:
             prompt = self.ps1
 
         if self.can_colorize:
+            t = THEME()
             prompt = (
-                f"{_colorize.theme["PROMPT"]}"
+                f"{t.prompt}"
                 f"{prompt}"
-                f"{_colorize.theme["RESET"]}"
+                f"{t.reset}"
             )
         return prompt
 
diff --git a/Lib/_pyrepl/utils.py b/Lib/_pyrepl/utils.py
@@ -23,6 +23,11 @@
 BUILTINS = {str(name) for name in dir(builtins) if not name.startswith('_')}
 
 
+def THEME():
+    # Not cached: the user can modify the theme inside the interactive session.
+    return _colorize.get_theme().repl
+
+
 class Span(NamedTuple):
     """Span indexing that's inclusive on both ends."""
 
@@ -44,7 +49,7 @@ def from_token(cls, token: TI, line_len: list[int]) -> Self:
 
 class ColorSpan(NamedTuple):
     span: Span
-    tag: _colorize.ColorTag
+    tag: str
 
 
 @functools.cache
@@ -135,7 +140,7 @@ def recover_unterminated_string(
 
         span = Span(start, end)
         trace("yielding span {a} -> {b}", a=span.start, b=span.end)
-        yield ColorSpan(span, "STRING")
+        yield ColorSpan(span, "string")
     else:
         trace(
             "unhandled token error({buffer}) = {te}",
@@ -164,28 +169,28 @@ def gen_colors_from_token_stream(
                 | T.TSTRING_START | T.TSTRING_MIDDLE | T.TSTRING_END
             ):
                 span = Span.from_token(token, line_lengths)
-                yield ColorSpan(span, "STRING")
+                yield ColorSpan(span, "string")
             case T.COMMENT:
                 span = Span.from_token(token, line_lengths)
-                yield ColorSpan(span, "COMMENT")
+                yield ColorSpan(span, "comment")
             case T.NUMBER:
                 span = Span.from_token(token, line_lengths)
-                yield ColorSpan(span, "NUMBER")
+                yield ColorSpan(span, "number")
             case T.OP:
                 if token.string in "([{":
                     bracket_level += 1
                 elif token.string in ")]}":
                     bracket_level -= 1
                 span = Span.from_token(token, line_lengths)
-                yield ColorSpan(span, "OP")
+                yield ColorSpan(span, "op")
             case T.NAME:
                 if is_def_name:
                     is_def_name = False
                     span = Span.from_token(token, line_lengths)
-                    yield ColorSpan(span, "DEFINITION")
+                    yield ColorSpan(span, "definition")
                 elif keyword.iskeyword(token.string):
                     span = Span.from_token(token, line_lengths)
-                    yield ColorSpan(span, "KEYWORD")
+                    yield ColorSpan(span, "keyword")
                     if token.string in IDENTIFIERS_AFTER:
                         is_def_name = True
                 elif (
@@ -194,10 +199,10 @@ def gen_colors_from_token_stream(
                     and is_soft_keyword_used(prev_token, token, next_token)
                 ):
                     span = Span.from_token(token, line_lengths)
-                    yield ColorSpan(span, "SOFT_KEYWORD")
+                    yield ColorSpan(span, "soft_keyword")
                 elif token.string in BUILTINS:
                     span = Span.from_token(token, line_lengths)
-                    yield ColorSpan(span, "BUILTIN")
+                    yield ColorSpan(span, "builtin")
 
 
 keyword_first_sets_match = {"False", "None", "True", "await", "lambda", "not"}
@@ -290,15 +295,16 @@ def disp_str(
         # move past irrelevant spans
         colors.pop(0)
 
+    theme = THEME()
     pre_color = ""
     post_color = ""
     if colors and colors[0].span.start < start_index:
         # looks like we're continuing a previous color (e.g. a multiline str)
-        pre_color = _colorize.theme[colors[0].tag]
+        pre_color = theme[colors[0].tag]
 
     for i, c in enumerate(buffer, start_index):
         if colors and colors[0].span.start == i:  # new color starts now
-            pre_color = _colorize.theme[colors[0].tag]
+            pre_color = theme[colors[0].tag]
 
         if c == "\x1a":  # CTRL-Z on Windows
             chars.append(c)
@@ -315,7 +321,7 @@ def disp_str(
             char_widths.append(str_width(c))
 
         if colors and colors[0].span.end == i:  # current color ends now
-            post_color = _colorize.theme["RESET"]
+            post_color = theme.reset
             colors.pop(0)
 
         chars[-1] = pre_color + chars[-1] + post_color
@@ -325,7 +331,7 @@ def disp_str(
     if colors and colors[0].span.start < i and colors[0].span.end > i:
         # even though the current color should be continued, reset it for now.
         # the next call to `disp_str()` will revive it.
-        chars[-1] += _colorize.theme["RESET"]
+        chars[-1] += theme.reset
 
     return chars, char_widths
 
diff --git a/Lib/test/test_pyrepl/test_reader.py b/Lib/test/test_pyrepl/test_reader.py
@@ -11,11 +11,11 @@
 from .support import reader_no_colors as prepare_reader
 from _pyrepl.console import Event
 from _pyrepl.reader import Reader
-from _colorize import theme
+from _colorize import get_theme
 
 
-overrides = {"RESET": "z", "SOFT_KEYWORD": "K"}
-colors = {overrides.get(k, k[0].lower()): v for k, v in theme.items()}
+overrides = {"reset": "z", "soft_keyword": "K"}
+colors = {overrides.get(k, k[0].lower()): v for k, v in get_theme().repl.items()}
 
 
 class TestReader(ScreenEqualMixin, TestCase):
