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

Co-authored-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com>

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
+    from typing import IO, Self, ClassVar
-
+    _theme: Theme
-    type ColorTag = Literal[
-        "PROMPT",
-        "KEYWORD",
-        "BUILTIN",
-        "COMMENT",
-        "STRING",
-        "NUMBER",
-        "OP",
-        "DEFINITION",
-        "SOFT_KEYWORD",
-        "RESET",
-    ]
-
-    theme: dict[ColorTag, str]
 
 
 class ANSIColors:
@@ -86,6 +75,186 @@ for attr, code in ANSIColors.__dict__.items():
         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:
+default_theme = Theme()
-    global theme
+theme_no_color = default_theme.no_colors()
-
-    if t:
-        theme = t
-        return
-
-    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,
-    }
 
 
-set_theme()
+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 not isinstance(t, Theme):
+        raise ValueError(f"Expected Theme object, found {t}")
+
+    _theme = t
+
+
+set_theme(default_theme)

Lib/_pyrepl/reader.py

@@ -28,7 +28,7 @@ from contextlib import contextmanager
 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 @@ class Reader:
             prompt = self.ps1
 
         if self.can_colorize:
-            prompt = (
+            t = THEME()
-                f"{_colorize.theme["PROMPT"]}"
+            prompt = f"{t.prompt}{prompt}{t.reset}"
-                f"{prompt}"
-                f"{_colorize.theme["RESET"]}"
-            )
         return prompt
 
     def push_input_trans(self, itrans: input.KeymapTranslator) -> None:

Lib/_pyrepl/utils.py

@@ -23,6 +23,11 @@ IDENTIFIERS_AFTER = {"def", "class"}
 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 @@ class Span(NamedTuple):
 
 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
 

Lib/argparse.py

@@ -176,13 +176,13 @@ class HelpFormatter(object):
             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 @@ class HelpFormatter(object):
 
             # 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,10 +312,7 @@ class HelpFormatter(object):
                         if part and part is not SUPPRESS])
 
     def _format_usage(self, usage, actions, groups, prefix):
-        bold_blue = self._ansi.BOLD_BLUE
+        t = self._theme
-        bold_magenta = self._ansi.BOLD_MAGENTA
-        magenta = self._ansi.MAGENTA
-        reset = self._ansi.RESET
 
         if prefix is None:
             prefix = _('usage: ')
@@ -325,15 +320,15 @@ class HelpFormatter(object):
         # 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}"}
+                % {"prog": f"{t.prog}{self._prog}{t.reset}{t.prog_extra}"}
-                + reset
+                + 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 @@ class HelpFormatter(object):
                 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 @@ class HelpFormatter(object):
 
         # collect all actions format strings
         parts = []
-        cyan = self._ansi.CYAN
+        t = self._theme
-        green = self._ansi.GREEN
-        yellow = self._ansi.YELLOW
-        reset = self._ansi.RESET
         for action in actions:
 
             # suppressed arguments are marked with None
@@ -465,7 +457,11 @@ class HelpFormatter(object):
             # 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_action
+                    + self._format_args(action, default)
+                    + t.reset
+                )
 
                 # if it's in a group, strip the outer []
                 if action in group_actions:
@@ -481,9 +477,9 @@ class HelpFormatter(object):
                 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
@@ -491,10 +487,13 @@ class HelpFormatter(object):
                     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}"
+                        option_color = t.summary_short_option
-                    part = f"{option_string} {yellow}{args_string}{reset}"
+                    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 @@ class HelpFormatter(object):
         return self._join_parts(parts)
 
     def _format_action_invocation(self, action):
-        bold_green = self._ansi.BOLD_GREEN
+        t = self._theme
-        bold_cyan = self._ansi.BOLD_CYAN
-        bold_yellow = self._ansi.BOLD_YELLOW
-        reset = self._ansi.RESET
 
         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 @@ class HelpFormatter(object):
                 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 @@ class HelpFormatter(object):
                 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
 

Lib/asyncio/__main__.py

@@ -12,7 +12,7 @@ import threading
 import types
 import warnings
 
-from _colorize import can_colorize, ANSIColors  # type: ignore[import-not-found]
+from _colorize import get_theme
 from _pyrepl.console import InteractiveColoredConsole
 
 from . import futures
@@ -103,8 +103,9 @@ class REPLThread(threading.Thread):
                     exec(startup_code, console.locals)
 
             ps1 = getattr(sys, "ps1", ">>> ")
-            if can_colorize() and CAN_USE_PYREPL:
+            if CAN_USE_PYREPL:
-                ps1 = f"{ANSIColors.BOLD_MAGENTA}{ps1}{ANSIColors.RESET}"
+                theme = get_theme().syntax
+                ps1 = f"{theme.prompt}{ps1}{theme.reset}"
             console.write(f"{ps1}import asyncio\n")
 
             if CAN_USE_PYREPL:

Lib/json/tool.py

@@ -7,7 +7,7 @@ import argparse
 import json
 import re
 import sys
-from _colorize import ANSIColors, can_colorize
+from _colorize import get_theme, can_colorize
 
 
 # The string we are colorizing is valid JSON,
@@ -17,27 +17,27 @@ from _colorize import ANSIColors, can_colorize
 _color_pattern = re.compile(r'''
     (?P<key>"(\\.|[^"\\])*")(?=:)           |
     (?P<string>"(\\.|[^"\\])*")             |
+    (?P<number>NaN|-?Infinity|[0-9\-+.Ee]+) |
     (?P<boolean>true|false)                 |
     (?P<null>null)
 ''', re.VERBOSE)
 
-
+_group_to_theme_color = {
-_colors = {
+    "key": "definition",
-    'key': ANSIColors.INTENSE_BLUE,
+    "string": "string",
-    'string': ANSIColors.BOLD_GREEN,
+    "number": "number",
-    'boolean': ANSIColors.BOLD_CYAN,
+    "boolean": "keyword",
-    'null': ANSIColors.BOLD_CYAN,
+    "null": "keyword",
 }
 
 
-def _replace_match_callback(match):
+def _colorize_json(json_str, theme):
-    for key, color in _colors.items():
+    def _replace_match_callback(match):
-        if m := match.group(key):
+        for group, color in _group_to_theme_color.items():
-            return f"{color}{m}{ANSIColors.RESET}"
+            if m := match.group(group):
-    return match.group()
+                return f"{theme[color]}{m}{theme.reset}"
+        return match.group()
 
-
-def _colorize_json(json_str):
     return re.sub(_color_pattern, _replace_match_callback, json_str)
 
 
@@ -100,13 +100,16 @@ def main():
         else:
             outfile = open(options.outfile, 'w', encoding='utf-8')
         with outfile:
-            for obj in objs:
+            if can_colorize(file=outfile):
-                if can_colorize(file=outfile):
+                t = get_theme(tty_file=outfile).syntax
+                for obj in objs:
                     json_str = json.dumps(obj, **dump_args)
-                    outfile.write(_colorize_json(json_str))
+                    outfile.write(_colorize_json(json_str, t))
-                else:
+                    outfile.write('\n')
+            else:
+                for obj in objs:
                     json.dump(obj, outfile, **dump_args)
-                outfile.write('\n')
+                    outfile.write('\n')
     except ValueError as e:
         raise SystemExit(e)
 

Lib/pdb.py

@@ -355,7 +355,7 @@ class Pdb(bdb.Bdb, cmd.Cmd):
         self._wait_for_mainpyfile = False
         self.tb_lineno = {}
         self.mode = mode
-        self.colorize = _colorize.can_colorize(file=stdout or sys.stdout) and colorize
+        self.colorize = colorize and _colorize.can_colorize(file=stdout or sys.stdout)
         # Try to load readline if it exists
         try:
             import readline

Lib/test/support/__init__.py

@@ -2855,36 +2855,59 @@ def iter_slot_wrappers(cls):
 
 
 @contextlib.contextmanager
-def no_color():
+def force_color(color: bool):
     import _colorize
     from .os_helper import EnvironmentVarGuard
 
     with (
-        swap_attr(_colorize, "can_colorize", lambda file=None: False),
+        swap_attr(_colorize, "can_colorize", lambda file=None: color),
         EnvironmentVarGuard() as env,
     ):
         env.unset("FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS")
-        env.set("NO_COLOR", "1")
+        env.set("FORCE_COLOR" if color else "NO_COLOR", "1")
         yield
 
 
-def force_not_colorized(func):
+def force_colorized(func):
-    """Force the terminal not to be colorized."""
+    """Force the terminal to be colorized."""
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
-        with no_color():
+        with force_color(True):
             return func(*args, **kwargs)
     return wrapper
 
 
-def force_not_colorized_test_class(cls):
+def force_not_colorized(func):
-    """Force the terminal not to be colorized for the entire test class."""
+    """Force the terminal NOT to be colorized."""
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        with force_color(False):
+            return func(*args, **kwargs)
+    return wrapper
+
+
+def force_colorized_test_class(cls):
+    """Force the terminal to be colorized for the entire test class."""
     original_setUpClass = cls.setUpClass
 
     @classmethod
     @functools.wraps(cls.setUpClass)
     def new_setUpClass(cls):
-        cls.enterClassContext(no_color())
+        cls.enterClassContext(force_color(True))
+        original_setUpClass()
+
+    cls.setUpClass = new_setUpClass
+    return cls
+
+
+def force_not_colorized_test_class(cls):
+    """Force the terminal NOT to be colorized for the entire test class."""
+    original_setUpClass = cls.setUpClass
+
+    @classmethod
+    @functools.wraps(cls.setUpClass)
+    def new_setUpClass(cls):
+        cls.enterClassContext(force_color(False))
         original_setUpClass()
 
     cls.setUpClass = new_setUpClass

Lib/test/test_argparse.py

@@ -7058,7 +7058,7 @@ class TestColorized(TestCase):
         super().setUp()
         # Ensure color even if ran with NO_COLOR=1
         _colorize.can_colorize = lambda *args, **kwargs: True
-        self.ansi = _colorize.ANSIColors()
+        self.theme = _colorize.get_theme(force_color=True).argparse
 
     def test_argparse_color(self):
         # Arrange: create a parser with a bit of everything
@@ -7120,13 +7120,17 @@ class TestColorized(TestCase):
         sub2 = subparsers.add_parser("sub2", deprecated=True, help="sub2 help")
         sub2.add_argument("--baz", choices=("X", "Y", "Z"), help="baz help")
 
-        heading = self.ansi.BOLD_BLUE
+        prog = self.theme.prog
-        label, label_b = self.ansi.YELLOW, self.ansi.BOLD_YELLOW
+        heading = self.theme.heading
-        long, long_b = self.ansi.CYAN, self.ansi.BOLD_CYAN
+        long = self.theme.summary_long_option
-        pos, pos_b = short, short_b = self.ansi.GREEN, self.ansi.BOLD_GREEN
+        short = self.theme.summary_short_option
-        sub = self.ansi.BOLD_GREEN
+        label = self.theme.summary_label
-        prog = self.ansi.BOLD_MAGENTA
+        pos = self.theme.summary_action
-        reset = self.ansi.RESET
+        long_b = self.theme.long_option
+        short_b = self.theme.short_option
+        label_b = self.theme.label
+        pos_b = self.theme.action
+        reset = self.theme.reset
 
         # Act
         help_text = parser.format_help()
@@ -7171,9 +7175,9 @@ class TestColorized(TestCase):
                 {heading}subcommands:{reset}
                   valid subcommands
 
-                  {sub}{{sub1,sub2}}{reset}           additional help
+                  {pos_b}{{sub1,sub2}}{reset}           additional help
-                    {sub}sub1{reset}                sub1 help
+                    {pos_b}sub1{reset}                sub1 help
-                    {sub}sub2{reset}                sub2 help
+                    {pos_b}sub2{reset}                sub2 help
                 """
             ),
         )
@@ -7187,10 +7191,10 @@ class TestColorized(TestCase):
             prog="PROG",
             usage="[prefix] %(prog)s [suffix]",
         )
-        heading = self.ansi.BOLD_BLUE
+        heading = self.theme.heading
-        prog = self.ansi.BOLD_MAGENTA
+        prog = self.theme.prog
-        reset = self.ansi.RESET
+        reset = self.theme.reset
-        usage = self.ansi.MAGENTA
+        usage = self.theme.prog_extra
 
         # Act
         help_text = parser.format_help()

Lib/test/test_json/test_tool.py

@@ -6,9 +6,11 @@ import unittest
 import subprocess
 
 from test import support
-from test.support import force_not_colorized, os_helper
+from test.support import force_colorized, force_not_colorized, os_helper
 from test.support.script_helper import assert_python_ok
 
+from _colorize import get_theme
+
 
 @support.requires_subprocess()
 class TestMain(unittest.TestCase):
@@ -246,34 +248,39 @@ class TestMain(unittest.TestCase):
         proc.communicate(b'"{}"')
         self.assertEqual(proc.returncode, errno.EPIPE)
 
+    @force_colorized
     def test_colors(self):
         infile = os_helper.TESTFN
         self.addCleanup(os.remove, infile)
 
+        t = get_theme().syntax
+        ob = "{"
+        cb = "}"
+
         cases = (
-            ('{}', b'{}'),
+            ('{}', '{}'),
-            ('[]', b'[]'),
+            ('[]', '[]'),
-            ('null', b'\x1b[1;36mnull\x1b[0m'),
+            ('null', f'{t.keyword}null{t.reset}'),
-            ('true', b'\x1b[1;36mtrue\x1b[0m'),
+            ('true', f'{t.keyword}true{t.reset}'),
-            ('false', b'\x1b[1;36mfalse\x1b[0m'),
+            ('false', f'{t.keyword}false{t.reset}'),
-            ('NaN', b'NaN'),
+            ('NaN', f'{t.number}NaN{t.reset}'),
-            ('Infinity', b'Infinity'),
+            ('Infinity', f'{t.number}Infinity{t.reset}'),
-            ('-Infinity', b'-Infinity'),
+            ('-Infinity', f'{t.number}-Infinity{t.reset}'),
-            ('"foo"', b'\x1b[1;32m"foo"\x1b[0m'),
+            ('"foo"', f'{t.string}"foo"{t.reset}'),
-            (r'" \"foo\" "', b'\x1b[1;32m" \\"foo\\" "\x1b[0m'),
+            (r'" \"foo\" "', f'{t.string}" \\"foo\\" "{t.reset}'),
-            ('"α"', b'\x1b[1;32m"\\u03b1"\x1b[0m'),
+            ('"α"', f'{t.string}"\\u03b1"{t.reset}'),
-            ('123', b'123'),
+            ('123', f'{t.number}123{t.reset}'),
-            ('-1.2345e+23', b'-1.2345e+23'),
+            ('-1.2345e+23', f'{t.number}-1.2345e+23{t.reset}'),
             (r'{"\\": ""}',
-             b'''\
+             f'''\
-{
+{ob}
-    \x1b[94m"\\\\"\x1b[0m: \x1b[1;32m""\x1b[0m
+    {t.definition}"\\\\"{t.reset}: {t.string}""{t.reset}
-}'''),
+{cb}'''),
             (r'{"\\\\": ""}',
-             b'''\
+             f'''\
-{
+{ob}
-    \x1b[94m"\\\\\\\\"\x1b[0m: \x1b[1;32m""\x1b[0m
+    {t.definition}"\\\\\\\\"{t.reset}: {t.string}""{t.reset}
-}'''),
+{cb}'''),
             ('''\
 {
     "foo": "bar",
@@ -281,30 +288,32 @@ class TestMain(unittest.TestCase):
     "qux": [true, false, null],
     "xyz": [NaN, -Infinity, Infinity]
 }''',
-             b'''\
+             f'''\
-{
+{ob}
-    \x1b[94m"foo"\x1b[0m: \x1b[1;32m"bar"\x1b[0m,
+    {t.definition}"foo"{t.reset}: {t.string}"bar"{t.reset},
-    \x1b[94m"baz"\x1b[0m: 1234,
+    {t.definition}"baz"{t.reset}: {t.number}1234{t.reset},
-    \x1b[94m"qux"\x1b[0m: [
+    {t.definition}"qux"{t.reset}: [
-        \x1b[1;36mtrue\x1b[0m,
+        {t.keyword}true{t.reset},
-        \x1b[1;36mfalse\x1b[0m,
+        {t.keyword}false{t.reset},
-        \x1b[1;36mnull\x1b[0m
+        {t.keyword}null{t.reset}
     ],
-    \x1b[94m"xyz"\x1b[0m: [
+    {t.definition}"xyz"{t.reset}: [
-        NaN,
+        {t.number}NaN{t.reset},
-        -Infinity,
+        {t.number}-Infinity{t.reset},
-        Infinity
+        {t.number}Infinity{t.reset}
     ]
-}'''),
+{cb}'''),
         )
 
         for input_, expected in cases:
             with self.subTest(input=input_):
                 with open(infile, "w", encoding="utf-8") as fp:
                     fp.write(input_)
-                _, stdout, _ = assert_python_ok('-m', self.module, infile,
+                _, stdout_b, _ = assert_python_ok(
-                                                PYTHON_COLORS='1')
+                    '-m', self.module, infile, FORCE_COLOR='1', __isolated='1'
-                stdout = stdout.replace(b'\r\n', b'\n')  # normalize line endings
+                )
+                stdout = stdout_b.decode()
+                stdout = stdout.replace('\r\n', '\n')  # normalize line endings
                 stdout = stdout.strip()
                 self.assertEqual(stdout, expected)
 

Lib/test/test_pdb.py

@@ -20,7 +20,7 @@ from asyncio.events import _set_event_loop_policy
 from contextlib import ExitStack, redirect_stdout
 from io import StringIO
 from test import support
-from test.support import force_not_colorized, has_socket_support, os_helper
+from test.support import has_socket_support, os_helper
 from test.support.import_helper import import_module
 from test.support.pty_helper import run_pty, FakeInput
 from test.support.script_helper import kill_python
@@ -3743,7 +3743,6 @@ def bœr():
         self.assertNotIn(b'Error', stdout,
                          "Got an error running test script under PDB")
 
-    @force_not_colorized
     def test_issue16180(self):
         # A syntax error in the debuggee.
         script = "def f: pass\n"
@@ -3757,7 +3756,6 @@ def bœr():
             'Fail to handle a syntax error in the debuggee.'
             .format(expected, stderr))
 
-    @force_not_colorized
     def test_issue84583(self):
         # A syntax error from ast.literal_eval should not make pdb exit.
         script = "import ast; ast.literal_eval('')\n"
@@ -4691,7 +4689,7 @@ class PdbTestInline(unittest.TestCase):
         self.assertIn("42", stdout)
 
 
-@unittest.skipUnless(_colorize.can_colorize(), "Test requires colorize")
+@support.force_colorized_test_class
 class PdbTestColorize(unittest.TestCase):
     def setUp(self):
         self._original_can_colorize = _colorize.can_colorize
@@ -4748,6 +4746,7 @@ class TestREPLSession(unittest.TestCase):
         self.assertEqual(p.returncode, 0)
 
 
+@support.force_not_colorized_test_class
 @support.requires_subprocess()
 class PdbTestReadline(unittest.TestCase):
     def setUpClass():

Lib/test/test_pyrepl/support.py

@@ -113,9 +113,6 @@ handle_events_narrow_console = partial(
     prepare_console=partial(prepare_console, width=10),
 )
 
-reader_no_colors = partial(prepare_reader, can_colorize=False)
-reader_force_colors = partial(prepare_reader, can_colorize=True)
-
 
 class FakeConsole(Console):
     def __init__(self, events, encoding="utf-8") -> None:

Lib/test/test_pyrepl/test_reader.py

@@ -4,20 +4,21 @@ import rlcompleter
 from textwrap import dedent
 from unittest import TestCase
 from unittest.mock import MagicMock
+from test.support import force_colorized_test_class, force_not_colorized_test_class
 
 from .support import handle_all_events, handle_events_narrow_console
 from .support import ScreenEqualMixin, code_to_events
-from .support import prepare_console, reader_force_colors
+from .support import prepare_reader, prepare_console
-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 default_theme
 
 
-overrides = {"RESET": "z", "SOFT_KEYWORD": "K"}
+overrides = {"reset": "z", "soft_keyword": "K"}
-colors = {overrides.get(k, k[0].lower()): v for k, v in theme.items()}
+colors = {overrides.get(k, k[0].lower()): v for k, v in default_theme.syntax.items()}
 
 
+@force_not_colorized_test_class
 class TestReader(ScreenEqualMixin, TestCase):
     def test_calc_screen_wrap_simple(self):
         events = code_to_events(10 * "a")
@@ -127,13 +128,6 @@ class TestReader(ScreenEqualMixin, TestCase):
         reader.setpos_from_xy(0, 0)
         self.assertEqual(reader.pos, 0)
 
-    def test_control_characters(self):
-        code = 'flag = "🏳️‍🌈"'
-        events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
-        self.assert_screen_equal(reader, 'flag = "🏳️\\u200d🌈"', clean=True)
-        self.assert_screen_equal(reader, 'flag {o}={z} {s}"🏳️\\u200d🌈"{z}'.format(**colors))
-
     def test_setpos_from_xy_multiple_lines(self):
         # fmt: off
         code = (
@@ -364,6 +358,8 @@ class TestReader(ScreenEqualMixin, TestCase):
         reader.setpos_from_xy(8, 0)
         self.assertEqual(reader.pos, 7)
 
+@force_colorized_test_class
+class TestReaderInColor(ScreenEqualMixin, TestCase):
     def test_syntax_highlighting_basic(self):
         code = dedent(
             """\
@@ -403,7 +399,7 @@ class TestReader(ScreenEqualMixin, TestCase):
         )
         expected_sync = expected.format(a="", **colors)
         events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(events)
         self.assert_screen_equal(reader, code, clean=True)
         self.assert_screen_equal(reader, expected_sync)
         self.assertEqual(reader.pos, 2**7 + 2**8)
@@ -416,7 +412,7 @@ class TestReader(ScreenEqualMixin, TestCase):
             [Event(evt="key", data="up", raw=bytearray(b"\x1bOA"))] * 13,
             code_to_events("async "),
         )
-        reader, _ = handle_all_events(more_events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(more_events)
         self.assert_screen_equal(reader, expected_async)
         self.assertEqual(reader.pos, 21)
         self.assertEqual(reader.cxy, (6, 1))
@@ -433,7 +429,7 @@ class TestReader(ScreenEqualMixin, TestCase):
             """
         ).format(**colors)
         events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(events)
         self.assert_screen_equal(reader, code, clean=True)
         self.assert_screen_equal(reader, expected)
 
@@ -451,7 +447,7 @@ class TestReader(ScreenEqualMixin, TestCase):
             """
         ).format(**colors)
         events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(events)
         self.assert_screen_equal(reader, code, clean=True)
         self.assert_screen_equal(reader, expected)
 
@@ -471,7 +467,7 @@ class TestReader(ScreenEqualMixin, TestCase):
             """
         ).format(**colors)
         events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(events)
         self.assert_screen_equal(reader, code, clean=True)
         self.assert_screen_equal(reader, expected)
 
@@ -497,6 +493,13 @@ class TestReader(ScreenEqualMixin, TestCase):
             """
         ).format(OB="{", CB="}", **colors)
         events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(events)
         self.assert_screen_equal(reader, code, clean=True)
         self.assert_screen_equal(reader, expected)
+
+    def test_control_characters(self):
+        code = 'flag = "🏳️‍🌈"'
+        events = code_to_events(code)
+        reader, _ = handle_all_events(events)
+        self.assert_screen_equal(reader, 'flag = "🏳️\\u200d🌈"', clean=True)
+        self.assert_screen_equal(reader, 'flag {o}={z} {s}"🏳️\\u200d🌈"{z}'.format(**colors))

Lib/test/test_pyrepl/test_unix_console.py

@@ -3,11 +3,12 @@ import os
 import sys
 import unittest
 from functools import partial
-from test.support import os_helper
+from test.support import os_helper, force_not_colorized_test_class
+
 from unittest import TestCase
 from unittest.mock import MagicMock, call, patch, ANY
 
-from .support import handle_all_events, code_to_events, reader_no_colors
+from .support import handle_all_events, code_to_events
 
 try:
     from _pyrepl.console import Event
@@ -33,12 +34,10 @@ def unix_console(events, **kwargs):
 
 handle_events_unix_console = partial(
     handle_all_events,
-    prepare_reader=reader_no_colors,
     prepare_console=unix_console,
 )
 handle_events_narrow_unix_console = partial(
     handle_all_events,
-    prepare_reader=reader_no_colors,
     prepare_console=partial(unix_console, width=5),
 )
 handle_events_short_unix_console = partial(
@@ -120,6 +119,7 @@ TERM_CAPABILITIES = {
 )
 @patch("termios.tcsetattr", lambda a, b, c: None)
 @patch("os.write")
+@force_not_colorized_test_class
 class TestConsole(TestCase):
     def test_simple_addition(self, _os_write):
         code = "12+34"
@@ -255,9 +255,7 @@ class TestConsole(TestCase):
         # fmt: on
 
         events = itertools.chain(code_to_events(code))
-        reader, console = handle_events_short_unix_console(
+        reader, console = handle_events_short_unix_console(events)
-            events, prepare_reader=reader_no_colors
-        )
 
         console.height = 2
         console.getheightwidth = MagicMock(lambda _: (2, 80))

Lib/test/test_pyrepl/test_windows_console.py

@@ -7,12 +7,13 @@ if sys.platform != "win32":
 
 import itertools
 from functools import partial
+from test.support import force_not_colorized_test_class
 from typing import Iterable
 from unittest import TestCase
 from unittest.mock import MagicMock, call
 
 from .support import handle_all_events, code_to_events
-from .support import reader_no_colors as default_prepare_reader
+from .support import prepare_reader as default_prepare_reader
 
 try:
     from _pyrepl.console import Event, Console
@@ -29,6 +30,7 @@ except ImportError:
     pass
 
 
+@force_not_colorized_test_class
 class WindowsConsoleTests(TestCase):
     def console(self, events, **kwargs) -> Console:
         console = WindowsConsole()

Lib/test/test_traceback.py

@@ -37,6 +37,12 @@ test_code.co_positions = lambda _: iter([(6, 6, 0, 0)])
 test_frame = namedtuple('frame', ['f_code', 'f_globals', 'f_locals'])
 test_tb = namedtuple('tb', ['tb_frame', 'tb_lineno', 'tb_next', 'tb_lasti'])
 
+color_overrides = {"reset": "z", "filename": "fn", "error_highlight": "E"}
+colors = {
+    color_overrides.get(k, k[0].lower()): v
+    for k, v in _colorize.default_theme.traceback.items()
+}
+
 
 LEVENSHTEIN_DATA_FILE = Path(__file__).parent / 'levenshtein_examples.json'
 
@@ -4721,6 +4727,8 @@ class MiscTest(unittest.TestCase):
 
 
 class TestColorizedTraceback(unittest.TestCase):
+    maxDiff = None
+
     def test_colorized_traceback(self):
         def foo(*args):
             x = {'a':{'b': None}}
@@ -4743,9 +4751,9 @@ class TestColorizedTraceback(unittest.TestCase):
                 e, capture_locals=True
             )
         lines = "".join(exc.format(colorize=True))
-        red = _colorize.ANSIColors.RED
+        red = colors["e"]
-        boldr = _colorize.ANSIColors.BOLD_RED
+        boldr = colors["E"]
-        reset = _colorize.ANSIColors.RESET
+        reset = colors["z"]
         self.assertIn("y = " + red + "x['a']['b']" + reset + boldr + "['c']" + reset, lines)
         self.assertIn("return " + red + "(lambda *args: foo(*args))" + reset + boldr + "(1,2,3,4)" + reset, lines)
         self.assertIn("return (lambda *args: " + red + "foo" + reset + boldr + "(*args)" + reset + ")(1,2,3,4)", lines)
@@ -4761,18 +4769,16 @@ class TestColorizedTraceback(unittest.TestCase):
                 e, capture_locals=True
             )
         actual = "".join(exc.format(colorize=True))
-        red = _colorize.ANSIColors.RED
+        def expected(t, m, fn, l, f, E, e, z):
-        magenta = _colorize.ANSIColors.MAGENTA
+            return "".join(
-        boldm = _colorize.ANSIColors.BOLD_MAGENTA
+                [
-        boldr = _colorize.ANSIColors.BOLD_RED
+                    f'  File {fn}"<string>"{z}, line {l}1{z}\n',
-        reset = _colorize.ANSIColors.RESET
+                    f'    a {E}${z} b\n',
-        expected = "".join([
+                    f'      {E}^{z}\n',
-        f'  File {magenta}"<string>"{reset}, line {magenta}1{reset}\n',
+                    f'{t}SyntaxError{z}: {m}invalid syntax{z}\n'
-        f'    a {boldr}${reset} b\n',
+                ]
-        f'      {boldr}^{reset}\n',
+            )
-        f'{boldm}SyntaxError{reset}: {magenta}invalid syntax{reset}\n']
+        self.assertIn(expected(**colors), actual)
-        )
-        self.assertIn(expected, actual)
 
     def test_colorized_traceback_is_the_default(self):
         def foo():
@@ -4788,23 +4794,21 @@ class TestColorizedTraceback(unittest.TestCase):
                     exception_print(e)
             actual = tbstderr.getvalue().splitlines()
 
-        red = _colorize.ANSIColors.RED
-        boldr = _colorize.ANSIColors.BOLD_RED
-        magenta = _colorize.ANSIColors.MAGENTA
-        boldm = _colorize.ANSIColors.BOLD_MAGENTA
-        reset = _colorize.ANSIColors.RESET
         lno_foo = foo.__code__.co_firstlineno
-        expected = ['Traceback (most recent call last):',
+        def expected(t, m, fn, l, f, E, e, z):
-            f'  File {magenta}"{__file__}"{reset}, '
+            return [
-            f'line {magenta}{lno_foo+5}{reset}, in {magenta}test_colorized_traceback_is_the_default{reset}',
+                'Traceback (most recent call last):',
-            f'    {red}foo{reset+boldr}(){reset}',
+                f'  File {fn}"{__file__}"{z}, '
-            f'    {red}~~~{reset+boldr}^^{reset}',
+                f'line {l}{lno_foo+5}{z}, in {f}test_colorized_traceback_is_the_default{z}',
-            f'  File {magenta}"{__file__}"{reset}, '
+                f'    {e}foo{z}{E}(){z}',
-            f'line {magenta}{lno_foo+1}{reset}, in {magenta}foo{reset}',
+                f'    {e}~~~{z}{E}^^{z}',
-            f'    {red}1{reset+boldr}/{reset+red}0{reset}',
+                f'  File {fn}"{__file__}"{z}, '
-            f'    {red}~{reset+boldr}^{reset+red}~{reset}',
+                f'line {l}{lno_foo+1}{z}, in {f}foo{z}',
-            f'{boldm}ZeroDivisionError{reset}: {magenta}division by zero{reset}']
+                f'    {e}1{z}{E}/{z}{e}0{z}',
-        self.assertEqual(actual, expected)
+                f'    {e}~{z}{E}^{z}{e}~{z}',
+                f'{t}ZeroDivisionError{z}: {m}division by zero{z}',
+            ]
+        self.assertEqual(actual, expected(**colors))
 
     def test_colorized_traceback_from_exception_group(self):
         def foo():
@@ -4822,33 +4826,31 @@ class TestColorizedTraceback(unittest.TestCase):
                 e, capture_locals=True
             )
 
-        red = _colorize.ANSIColors.RED
-        boldr = _colorize.ANSIColors.BOLD_RED
-        magenta = _colorize.ANSIColors.MAGENTA
-        boldm = _colorize.ANSIColors.BOLD_MAGENTA
-        reset = _colorize.ANSIColors.RESET
         lno_foo = foo.__code__.co_firstlineno
         actual = "".join(exc.format(colorize=True)).splitlines()
-        expected = [f"  + Exception Group Traceback (most recent call last):",
+        def expected(t, m, fn, l, f, E, e, z):
-                   f'  |   File {magenta}"{__file__}"{reset}, line {magenta}{lno_foo+9}{reset}, in {magenta}test_colorized_traceback_from_exception_group{reset}',
+            return [
-                   f'  |     {red}foo{reset}{boldr}(){reset}',
+                f"  + Exception Group Traceback (most recent call last):",
-                   f'  |     {red}~~~{reset}{boldr}^^{reset}',
+                f'  |   File {fn}"{__file__}"{z}, line {l}{lno_foo+9}{z}, in {f}test_colorized_traceback_from_exception_group{z}',
-                   f"  |     e = ExceptionGroup('test', [ZeroDivisionError('division by zero')])",
+                f'  |     {e}foo{z}{E}(){z}',
-                   f"  |     foo = {foo}",
+                f'  |     {e}~~~{z}{E}^^{z}',
-                   f'  |     self = <{__name__}.TestColorizedTraceback testMethod=test_colorized_traceback_from_exception_group>',
+                f"  |     e = ExceptionGroup('test', [ZeroDivisionError('division by zero')])",
-                   f'  |   File {magenta}"{__file__}"{reset}, line {magenta}{lno_foo+6}{reset}, in {magenta}foo{reset}',
+                f"  |     foo = {foo}",
-                   f'  |     raise ExceptionGroup("test", exceptions)',
+                f'  |     self = <{__name__}.TestColorizedTraceback testMethod=test_colorized_traceback_from_exception_group>',
-                   f"  |     exceptions = [ZeroDivisionError('division by zero')]",
+                f'  |   File {fn}"{__file__}"{z}, line {l}{lno_foo+6}{z}, in {f}foo{z}',
-                   f'  | {boldm}ExceptionGroup{reset}: {magenta}test (1 sub-exception){reset}',
+                f'  |     raise ExceptionGroup("test", exceptions)',
-                   f'  +-+---------------- 1 ----------------',
+                f"  |     exceptions = [ZeroDivisionError('division by zero')]",
-                   f'    | Traceback (most recent call last):',
+                f'  | {t}ExceptionGroup{z}: {m}test (1 sub-exception){z}',
-                   f'    |   File {magenta}"{__file__}"{reset}, line {magenta}{lno_foo+3}{reset}, in {magenta}foo{reset}',
+                f'  +-+---------------- 1 ----------------',
-                   f'    |     {red}1 {reset}{boldr}/{reset}{red} 0{reset}',
+                f'    | Traceback (most recent call last):',
-                   f'    |     {red}~~{reset}{boldr}^{reset}{red}~~{reset}',
+                f'    |   File {fn}"{__file__}"{z}, line {l}{lno_foo+3}{z}, in {f}foo{z}',
-                   f"    |     exceptions = [ZeroDivisionError('division by zero')]",
+                f'    |     {e}1 {z}{E}/{z}{e} 0{z}',
-                   f'    | {boldm}ZeroDivisionError{reset}: {magenta}division by zero{reset}',
+                f'    |     {e}~~{z}{E}^{z}{e}~~{z}',
-                   f'    +------------------------------------']
+                f"    |     exceptions = [ZeroDivisionError('division by zero')]",
-        self.assertEqual(actual, expected)
+                f'    | {t}ZeroDivisionError{z}: {m}division by zero{z}',
+                f'    +------------------------------------',
+        ]
+        self.assertEqual(actual, expected(**colors))
 
 if __name__ == "__main__":
     unittest.main()

Lib/traceback.py

@@ -10,9 +10,9 @@ import codeop
 import keyword
 import tokenize
 import io
-from contextlib import suppress
 import _colorize
-from _colorize import ANSIColors
+
+from contextlib import suppress
 
 __all__ = ['extract_stack', 'extract_tb', 'format_exception',
            'format_exception_only', 'format_list', 'format_stack',
@@ -187,15 +187,13 @@ def _format_final_exc_line(etype, value, *, insert_final_newline=True, colorize=
     valuestr = _safe_string(value, 'exception')
     end_char = "\n" if insert_final_newline else ""
     if colorize:
-        if value is None or not valuestr:
+        theme = _colorize.get_theme(force_color=True).traceback
-            line = f"{ANSIColors.BOLD_MAGENTA}{etype}{ANSIColors.RESET}{end_char}"
-        else:
-            line = f"{ANSIColors.BOLD_MAGENTA}{etype}{ANSIColors.RESET}: {ANSIColors.MAGENTA}{valuestr}{ANSIColors.RESET}{end_char}"
     else:
-        if value is None or not valuestr:
+        theme = _colorize.get_theme(force_no_color=True).traceback
-            line = f"{etype}{end_char}"
+    if value is None or not valuestr:
-        else:
+        line = f"{theme.type}{etype}{theme.reset}{end_char}"
-            line = f"{etype}: {valuestr}{end_char}"
+    else:
+        line = f"{theme.type}{etype}{theme.reset}: {theme.message}{valuestr}{theme.reset}{end_char}"
     return line
 
 
@@ -539,21 +537,22 @@ class StackSummary(list):
         if frame_summary.filename.startswith("<stdin>-"):
             filename = "<stdin>"
         if colorize:
-            row.append('  File {}"{}"{}, line {}{}{}, in {}{}{}\n'.format(
+            theme = _colorize.get_theme(force_color=True).traceback
-                    ANSIColors.MAGENTA,
-                    filename,
-                    ANSIColors.RESET,
-                    ANSIColors.MAGENTA,
-                    frame_summary.lineno,
-                    ANSIColors.RESET,
-                    ANSIColors.MAGENTA,
-                    frame_summary.name,
-                    ANSIColors.RESET,
-                    )
-            )
         else:
-            row.append('  File "{}", line {}, in {}\n'.format(
+            theme = _colorize.get_theme(force_no_color=True).traceback
-                filename, frame_summary.lineno, frame_summary.name))
+        row.append(
+            '  File {}"{}"{}, line {}{}{}, in {}{}{}\n'.format(
+                theme.filename,
+                filename,
+                theme.reset,
+                theme.line_no,
+                frame_summary.lineno,
+                theme.reset,
+                theme.frame,
+                frame_summary.name,
+                theme.reset,
+            )
+        )
         if frame_summary._dedented_lines and frame_summary._dedented_lines.strip():
             if (
                 frame_summary.colno is None or
@@ -672,11 +671,11 @@ class StackSummary(list):
                         for color, group in itertools.groupby(itertools.zip_longest(line, carets, fillvalue=""), key=lambda x: x[1]):
                             caret_group = list(group)
                             if color == "^":
-                                colorized_line_parts.append(ANSIColors.BOLD_RED + "".join(char for char, _ in caret_group) + ANSIColors.RESET)
+                                colorized_line_parts.append(theme.error_highlight + "".join(char for char, _ in caret_group) + theme.reset)
-                                colorized_carets_parts.append(ANSIColors.BOLD_RED + "".join(caret for _, caret in caret_group) + ANSIColors.RESET)
+                                colorized_carets_parts.append(theme.error_highlight + "".join(caret for _, caret in caret_group) + theme.reset)
                             elif color == "~":
-                                colorized_line_parts.append(ANSIColors.RED + "".join(char for char, _ in caret_group) + ANSIColors.RESET)
+                                colorized_line_parts.append(theme.error_range + "".join(char for char, _ in caret_group) + theme.reset)
-                                colorized_carets_parts.append(ANSIColors.RED + "".join(caret for _, caret in caret_group) + ANSIColors.RESET)
+                                colorized_carets_parts.append(theme.error_range + "".join(caret for _, caret in caret_group) + theme.reset)
                             else:
                                 colorized_line_parts.append("".join(char for char, _ in caret_group))
                                 colorized_carets_parts.append("".join(caret for _, caret in caret_group))
@@ -1378,20 +1377,20 @@ class TracebackException:
         """Format SyntaxError exceptions (internal helper)."""
         # Show exactly where the problem was found.
         colorize = kwargs.get("colorize", False)
+        if colorize:
+            theme = _colorize.get_theme(force_color=True).traceback
+        else:
+            theme = _colorize.get_theme(force_no_color=True).traceback
         filename_suffix = ''
         if self.lineno is not None:
-            if colorize:
+            yield '  File {}"{}"{}, line {}{}{}\n'.format(
-                yield '  File {}"{}"{}, line {}{}{}\n'.format(
+                theme.filename,
-                    ANSIColors.MAGENTA,
+                self.filename or "<string>",
-                    self.filename or "<string>",
+                theme.reset,
-                    ANSIColors.RESET,
+                theme.line_no,
-                    ANSIColors.MAGENTA,
+                self.lineno,
-                    self.lineno,
+                theme.reset,
-                    ANSIColors.RESET,
+                )
-                    )
-            else:
-                yield '  File "{}", line {}\n'.format(
-                    self.filename or "<string>", self.lineno)
         elif self.filename is not None:
             filename_suffix = ' ({})'.format(self.filename)
 
@@ -1441,11 +1440,11 @@ class TracebackException:
                         # colorize from colno to end_colno
                         ltext = (
                             ltext[:colno] +
-                            ANSIColors.BOLD_RED + ltext[colno:end_colno] + ANSIColors.RESET +
+                            theme.error_highlight + ltext[colno:end_colno] + theme.reset +
                             ltext[end_colno:]
                         )
-                        start_color = ANSIColors.BOLD_RED
+                        start_color = theme.error_highlight
-                        end_color = ANSIColors.RESET
+                        end_color = theme.reset
                     yield '    {}\n'.format(ltext)
                     yield '    {}{}{}{}\n'.format(
                         "".join(caretspace),
@@ -1456,17 +1455,15 @@ class TracebackException:
                 else:
                     yield '    {}\n'.format(ltext)
         msg = self.msg or "<no detail available>"
-        if colorize:
+        yield "{}{}{}: {}{}{}{}\n".format(
-            yield "{}{}{}: {}{}{}{}\n".format(
+            theme.type,
-                ANSIColors.BOLD_MAGENTA,
+            stype,
-                stype,
+            theme.reset,
-                ANSIColors.RESET,
+            theme.message,
-                ANSIColors.MAGENTA,
+            msg,
-                msg,
+            theme.reset,
-                ANSIColors.RESET,
+            filename_suffix,
-                filename_suffix)
+        )
-        else:
-            yield "{}: {}{}\n".format(stype, msg, filename_suffix)
 
     def format(self, *, chain=True, _ctx=None, **kwargs):
         """Format the exception.

Lib/unittest/runner.py

@@ -4,7 +4,7 @@ import sys
 import time
 import warnings
 
-from _colorize import get_colors
+from _colorize import get_theme
 
 from . import result
 from .case import _SubTest
@@ -45,7 +45,7 @@ class TextTestResult(result.TestResult):
         self.showAll = verbosity > 1
         self.dots = verbosity == 1
         self.descriptions = descriptions
-        self._ansi = get_colors(file=stream)
+        self._theme = get_theme(tty_file=stream).unittest
         self._newline = True
         self.durations = durations
 
@@ -79,101 +79,99 @@ class TextTestResult(result.TestResult):
 
     def addSubTest(self, test, subtest, err):
         if err is not None:
-            red, reset = self._ansi.RED, self._ansi.RESET
+            t = self._theme
             if self.showAll:
                 if issubclass(err[0], subtest.failureException):
-                    self._write_status(subtest, f"{red}FAIL{reset}")
+                    self._write_status(subtest, f"{t.fail}FAIL{t.reset}")
                 else:
-                    self._write_status(subtest, f"{red}ERROR{reset}")
+                    self._write_status(subtest, f"{t.fail}ERROR{t.reset}")
             elif self.dots:
                 if issubclass(err[0], subtest.failureException):
-                    self.stream.write(f"{red}F{reset}")
+                    self.stream.write(f"{t.fail}F{t.reset}")
                 else:
-                    self.stream.write(f"{red}E{reset}")
+                    self.stream.write(f"{t.fail}E{t.reset}")
                 self.stream.flush()
         super(TextTestResult, self).addSubTest(test, subtest, err)
 
     def addSuccess(self, test):
         super(TextTestResult, self).addSuccess(test)
-        green, reset = self._ansi.GREEN, self._ansi.RESET
+        t = self._theme
         if self.showAll:
-            self._write_status(test, f"{green}ok{reset}")
+            self._write_status(test, f"{t.passed}ok{t.reset}")
         elif self.dots:
-            self.stream.write(f"{green}.{reset}")
+            self.stream.write(f"{t.passed}.{t.reset}")
             self.stream.flush()
 
     def addError(self, test, err):
         super(TextTestResult, self).addError(test, err)
-        red, reset = self._ansi.RED, self._ansi.RESET
+        t = self._theme
         if self.showAll:
-            self._write_status(test, f"{red}ERROR{reset}")
+            self._write_status(test, f"{t.fail}ERROR{t.reset}")
         elif self.dots:
-            self.stream.write(f"{red}E{reset}")
+            self.stream.write(f"{t.fail}E{t.reset}")
             self.stream.flush()
 
     def addFailure(self, test, err):
         super(TextTestResult, self).addFailure(test, err)
-        red, reset = self._ansi.RED, self._ansi.RESET
+        t = self._theme
         if self.showAll:
-            self._write_status(test, f"{red}FAIL{reset}")
+            self._write_status(test, f"{t.fail}FAIL{t.reset}")
         elif self.dots:
-            self.stream.write(f"{red}F{reset}")
+            self.stream.write(f"{t.fail}F{t.reset}")
             self.stream.flush()
 
     def addSkip(self, test, reason):
         super(TextTestResult, self).addSkip(test, reason)
-        yellow, reset = self._ansi.YELLOW, self._ansi.RESET
+        t = self._theme
         if self.showAll:
-            self._write_status(test, f"{yellow}skipped{reset} {reason!r}")
+            self._write_status(test, f"{t.warn}skipped{t.reset} {reason!r}")
         elif self.dots:
-            self.stream.write(f"{yellow}s{reset}")
+            self.stream.write(f"{t.warn}s{t.reset}")
             self.stream.flush()
 
     def addExpectedFailure(self, test, err):
         super(TextTestResult, self).addExpectedFailure(test, err)
-        yellow, reset = self._ansi.YELLOW, self._ansi.RESET
+        t = self._theme
         if self.showAll:
-            self.stream.writeln(f"{yellow}expected failure{reset}")
+            self.stream.writeln(f"{t.warn}expected failure{t.reset}")
             self.stream.flush()
         elif self.dots:
-            self.stream.write(f"{yellow}x{reset}")
+            self.stream.write(f"{t.warn}x{t.reset}")
             self.stream.flush()
 
     def addUnexpectedSuccess(self, test):
         super(TextTestResult, self).addUnexpectedSuccess(test)
-        red, reset = self._ansi.RED, self._ansi.RESET
+        t = self._theme
         if self.showAll:
-            self.stream.writeln(f"{red}unexpected success{reset}")
+            self.stream.writeln(f"{t.fail}unexpected success{t.reset}")
             self.stream.flush()
         elif self.dots:
-            self.stream.write(f"{red}u{reset}")
+            self.stream.write(f"{t.fail}u{t.reset}")
             self.stream.flush()
 
     def printErrors(self):
-        bold_red = self._ansi.BOLD_RED
+        t = self._theme
-        red = self._ansi.RED
-        reset = self._ansi.RESET
         if self.dots or self.showAll:
             self.stream.writeln()
             self.stream.flush()
-        self.printErrorList(f"{red}ERROR{reset}", self.errors)
+        self.printErrorList(f"{t.fail}ERROR{t.reset}", self.errors)
-        self.printErrorList(f"{red}FAIL{reset}", self.failures)
+        self.printErrorList(f"{t.fail}FAIL{t.reset}", self.failures)
         unexpectedSuccesses = getattr(self, "unexpectedSuccesses", ())
         if unexpectedSuccesses:
             self.stream.writeln(self.separator1)
             for test in unexpectedSuccesses:
                 self.stream.writeln(
-                    f"{red}UNEXPECTED SUCCESS{bold_red}: "
+                    f"{t.fail}UNEXPECTED SUCCESS{t.fail_info}: "
-                    f"{self.getDescription(test)}{reset}"
+                    f"{self.getDescription(test)}{t.reset}"
                 )
             self.stream.flush()
 
     def printErrorList(self, flavour, errors):
-        bold_red, reset = self._ansi.BOLD_RED, self._ansi.RESET
+        t = self._theme
         for test, err in errors:
             self.stream.writeln(self.separator1)
             self.stream.writeln(
-                f"{flavour}{bold_red}: {self.getDescription(test)}{reset}"
+                f"{flavour}{t.fail_info}: {self.getDescription(test)}{t.reset}"
             )
             self.stream.writeln(self.separator2)
             self.stream.writeln("%s" % err)
@@ -286,31 +284,26 @@ class TextTestRunner(object):
             expected_fails, unexpected_successes, skipped = results
 
         infos = []
-        ansi = get_colors(file=self.stream)
+        t = get_theme(tty_file=self.stream).unittest
-        bold_red = ansi.BOLD_RED
-        green = ansi.GREEN
-        red = ansi.RED
-        reset = ansi.RESET
-        yellow = ansi.YELLOW
 
         if not result.wasSuccessful():
-            self.stream.write(f"{bold_red}FAILED{reset}")
+            self.stream.write(f"{t.fail_info}FAILED{t.reset}")
             failed, errored = len(result.failures), len(result.errors)
             if failed:
-                infos.append(f"{bold_red}failures={failed}{reset}")
+                infos.append(f"{t.fail_info}failures={failed}{t.reset}")
             if errored:
-                infos.append(f"{bold_red}errors={errored}{reset}")
+                infos.append(f"{t.fail_info}errors={errored}{t.reset}")
         elif run == 0 and not skipped:
-            self.stream.write(f"{yellow}NO TESTS RAN{reset}")
+            self.stream.write(f"{t.warn}NO TESTS RAN{t.reset}")
         else:
-            self.stream.write(f"{green}OK{reset}")
+            self.stream.write(f"{t.passed}OK{t.reset}")
         if skipped:
-            infos.append(f"{yellow}skipped={skipped}{reset}")
+            infos.append(f"{t.warn}skipped={skipped}{t.reset}")
         if expected_fails:
-            infos.append(f"{yellow}expected failures={expected_fails}{reset}")
+            infos.append(f"{t.warn}expected failures={expected_fails}{t.reset}")
         if unexpected_successes:
             infos.append(
-                f"{red}unexpected successes={unexpected_successes}{reset}"
+                f"{t.fail}unexpected successes={unexpected_successes}{t.reset}"
             )
         if infos:
             self.stream.writeln(" (%s)" % (", ".join(infos),))

Misc/NEWS.d/next/Core_and_Builtins/2025-05-04-19-46-14.gh-issue-133346.nRXi4f.rst

@@ -0,0 +1 @@
+Added experimental color theming support to the ``_colorize`` module.