diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py index 173b270f06a170..db29f0428a19ba 100644 --- a/Lib/test/support/__init__.py +++ b/Lib/test/support/__init__.py @@ -17,6 +17,7 @@ import types import unittest import warnings +from collections.abc import Callable __all__ = [ @@ -2833,31 +2834,45 @@ def is_slot_wrapper(name, value): yield name, True +def _disable_terminal_color() -> Callable[[], bool]: + import _colorize + + original_fn = _colorize.can_colorize + variables: dict[str, str | None] = { + "PYTHON_COLORS": None, + "FORCE_COLOR": None, + "NO_COLOR": None, + } + for key in variables: + variables[key] = os.environ.pop(key, None) + os.environ["NO_COLOR"] = "1" + _colorize.can_colorize = lambda: False + return original_fn, variables + + +def _re_enable_terminal_color( + original_fn: Callable[[], bool], variables: dict[str, str | None] +): + import _colorize + + _colorize.can_colorize = original_fn + del os.environ["NO_COLOR"] + for key, value in variables.items(): + if value is not None: + os.environ[key] = value + + def force_not_colorized(func): """Force the terminal not to be colorized.""" @functools.wraps(func) def wrapper(*args, **kwargs): - import _colorize - original_fn = _colorize.can_colorize - variables: dict[str, str | None] = { - "PYTHON_COLORS": None, "FORCE_COLOR": None, "NO_COLOR": None - } try: - for key in variables: - variables[key] = os.environ.pop(key, None) - os.environ["NO_COLOR"] = "1" - _colorize.can_colorize = lambda: False + original_fn, variables = _disable_terminal_color() return func(*args, **kwargs) finally: - _colorize.can_colorize = original_fn - del os.environ["NO_COLOR"] - for key, value in variables.items(): - if value is not None: - os.environ[key] = value + _re_enable_terminal_color(original_fn, variables) return wrapper - - def force_not_colorized_test_class(cls): """Force the terminal not to be colorized.""" original_setup = cls.setUp @@ -2865,29 +2880,13 @@ def force_not_colorized_test_class(cls): @functools.wraps(cls.setUp) def setUp_wrapper(self, *args, **kwargs): - import _colorize + self._original_fn, self._variables = _disable_terminal_color() - self._original_fn = _colorize.can_colorize - self._variables: dict[str, str | None] = { - "PYTHON_COLORS": None, - "FORCE_COLOR": None, - "NO_COLOR": None, - } - for key in self._variables: - self._variables[key] = os.environ.pop(key, None) - os.environ["NO_COLOR"] = "1" - _colorize.can_colorize = lambda: False return original_setup(self, *args, **kwargs) @functools.wraps(cls.tearDown) def tearDown_wrapper(self, *args, **kwargs): - import _colorize - - _colorize.can_colorize = self._original_fn - del os.environ["NO_COLOR"] - for key, value in self._variables.items(): - if value is not None: - os.environ[key] = value + _re_enable_terminal_color(self._original_fn, self._variables) return original_teardown(self, *args, **kwargs) cls.setUp = setUp_wrapper