diff --git a/Lib/_colorize.py b/Lib/_colorize.py index 9eb6f0933..d6673f669 100644 --- a/Lib/_colorize.py +++ b/Lib/_colorize.py @@ -1,13 +1,16 @@ -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 + from typing import IO, Self, ClassVar + _theme: Theme class ANSIColors: @@ -17,11 +20,13 @@ class ANSIColors: BLUE = "\x1b[34m" CYAN = "\x1b[36m" GREEN = "\x1b[32m" + GREY = "\x1b[90m" MAGENTA = "\x1b[35m" RED = "\x1b[31m" WHITE = "\x1b[37m" # more like LIGHT GRAY YELLOW = "\x1b[33m" + BOLD = "\x1b[1m" BOLD_BLACK = "\x1b[1;30m" # DARK GRAY BOLD_BLUE = "\x1b[1;34m" BOLD_CYAN = "\x1b[1;36m" @@ -60,13 +65,196 @@ class ANSIColors: INTENSE_BACKGROUND_YELLOW = "\x1b[103m" +ColorCodes = set() NoColors = ANSIColors() -for attr in dir(NoColors): +for attr, code in ANSIColors.__dict__.items(): if not attr.startswith("__"): + ColorCodes.add(code) 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, kw_only=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 + keyword_constant: 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: @@ -76,22 +264,37 @@ def get_colors( return NoColors +def decolor(text: str) -> str: + """Remove ANSI color codes from a string.""" + for code in ColorCodes: + text = text.replace(code, "") + return text + + def can_colorize(*, file: IO[str] | IO[bytes] | None = None) -> bool: + + def _safe_getenv(k: str, fallback: str | None = None) -> str | None: + """Exception-safe environment retrieval. See gh-128636.""" + try: + return os.environ.get(k, fallback) + except Exception: + return fallback + if file is None: file = sys.stdout if not sys.flags.ignore_environment: - if os.environ.get("PYTHON_COLORS") == "0": + if _safe_getenv("PYTHON_COLORS") == "0": return False - if os.environ.get("PYTHON_COLORS") == "1": + if _safe_getenv("PYTHON_COLORS") == "1": return True - if os.environ.get("NO_COLOR"): + if _safe_getenv("NO_COLOR"): return False if not COLORIZE: return False - if os.environ.get("FORCE_COLOR"): + if _safe_getenv("FORCE_COLOR"): return True - if os.environ.get("TERM") == "dumb": + if _safe_getenv("TERM") == "dumb": return False if not hasattr(file, "fileno"): @@ -108,5 +311,45 @@ def can_colorize(*, file: IO[str] | IO[bytes] | None = None) -> bool: try: return os.isatty(file.fileno()) - except io.UnsupportedOperation: + except OSError: return hasattr(file, "isatty") and file.isatty() + + +default_theme = Theme() +theme_no_color = default_theme.no_colors() + + +def get_theme( + *, + tty_file: IO[str] | IO[bytes] | None = None, + force_color: bool = False, + force_no_color: bool = False, +) -> Theme: + """Returns the currently set theme, potentially in a zero-color variant. + + In cases where colorizing is not possible (see `can_colorize`), the returned + theme contains all empty strings in all color definitions. + See `Theme.no_colors()` for more information. + + It is recommended not to cache the result of this function for extended + periods of time because the user might influence theme selection by + the interactive shell, a debugger, or application-specific code. The + environment (including environment variable state and console configuration + on Windows) can also change in the course of the application life cycle. + """ + if force_color or (not force_no_color and + can_colorize(file=tty_file)): + return _theme + return theme_no_color + + +def set_theme(t: Theme) -> None: + global _theme + + if not isinstance(t, Theme): + raise ValueError(f"Expected Theme object, found {t}") + + _theme = t + + +set_theme(default_theme) diff --git a/Lib/test/test__colorize.py b/Lib/test/test__colorize.py index b2f0bb138..31dc60fec 100644 --- a/Lib/test/test__colorize.py +++ b/Lib/test/test__colorize.py @@ -1,4 +1,5 @@ import contextlib +import dataclasses import io import sys import unittest @@ -21,6 +22,42 @@ def supports_virtual_terminal(): return contextlib.nullcontext() +class TestTheme(unittest.TestCase): + + @unittest.expectedFailure # TODO: RUSTPYTHON + def test_attributes(self): + # only theme configurations attributes by default + for field in dataclasses.fields(_colorize.Theme): + with self.subTest(field.name): + self.assertIsSubclass(field.type, _colorize.ThemeSection) + self.assertIsNotNone(field.default_factory) + + def test_copy_with(self): + theme = _colorize.Theme() + + copy = theme.copy_with() + self.assertEqual(theme, copy) + + unittest_no_colors = _colorize.Unittest.no_colors() + copy = theme.copy_with(unittest=unittest_no_colors) + self.assertEqual(copy.argparse, theme.argparse) + self.assertEqual(copy.syntax, theme.syntax) + self.assertEqual(copy.traceback, theme.traceback) + self.assertEqual(copy.unittest, unittest_no_colors) + + def test_no_colors(self): + # idempotence test + theme_no_colors = _colorize.Theme().no_colors() + theme_no_colors_no_colors = theme_no_colors.no_colors() + self.assertEqual(theme_no_colors, theme_no_colors_no_colors) + + # attributes check + for section in dataclasses.fields(_colorize.Theme): + with self.subTest(section.name): + section_theme = getattr(theme_no_colors, section.name) + self.assertEqual(section_theme, section.type.no_colors()) + + class TestColorizeFunction(unittest.TestCase): def test_colorized_detection_checks_for_environment_variables(self): def check(env, fallback, expected): @@ -129,6 +166,17 @@ class TestColorizeFunction(unittest.TestCase): file.isatty.return_value = False self.assertEqual(_colorize.can_colorize(file=file), False) + # The documentation for file.fileno says: + # > An OSError is raised if the IO object does not use a file descriptor. + # gh-141570: Check OSError is caught and handled + with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError): + file = unittest.mock.MagicMock() + file.fileno.side_effect = OSError + file.isatty.return_value = True + self.assertEqual(_colorize.can_colorize(file=file), True) + file.isatty.return_value = False + self.assertEqual(_colorize.can_colorize(file=file), False) + if __name__ == "__main__": unittest.main()