python · ambv · May 5, 2025 · May 3, 2025 · May 3, 2025 · May 3, 2025
@@ -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,165 @@ 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.
+
+
+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.
+    """
+    __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)
+
+    def no_colors(self) -> Self:
+        color_state: dict[str, str] = {}
+        for color_name in self.__dataclass_fields__:
+            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 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,
+        )
+
+    def no_colors(self) -> 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 type(self)(
+            argparse=self.argparse.no_colors(),
+            syntax=self.syntax.no_colors(),
+            traceback=self.traceback.no_colors(),
+            unittest=self.unittest.no_colors(),
+        )
+
+
 def get_colors(
     colorize: bool = False, *, file: IO[str] | IO[bytes] | None = None
 ) -> ANSIColors:
@@ -138,26 +286,37 @@ 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
+
 
-    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
 
 
-set_theme()
+set_theme(default_theme)
@@ -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:

@@ -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().syntax
+
+
 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