Skip to content

Commit

Permalink
pythongh-117225: Move colorize functionality to own internal module (p…
Browse files Browse the repository at this point in the history
  • Loading branch information
hugovk authored and SonicField committed May 8, 2024
1 parent 7e10b3f commit 1e0a640
Show file tree
Hide file tree
Showing 8 changed files with 218 additions and 180 deletions.
64 changes: 64 additions & 0 deletions Lib/_colorize.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
import io
import os
import sys

COLORIZE = True


class ANSIColors:
BOLD_GREEN = "\x1b[1;32m"
BOLD_MAGENTA = "\x1b[1;35m"
BOLD_RED = "\x1b[1;31m"
GREEN = "\x1b[32m"
GREY = "\x1b[90m"
MAGENTA = "\x1b[35m"
RED = "\x1b[31m"
RESET = "\x1b[0m"
YELLOW = "\x1b[33m"


NoColors = ANSIColors()

for attr in dir(NoColors):
if not attr.startswith("__"):
setattr(NoColors, attr, "")


def get_colors(colorize: bool = False) -> ANSIColors:
if colorize or can_colorize():
return ANSIColors()
else:
return NoColors


def can_colorize() -> bool:
if sys.platform == "win32":
try:
import nt

if not nt._supports_virtual_terminal():
return False
except (ImportError, AttributeError):
return False
if not sys.flags.ignore_environment:
if os.environ.get("PYTHON_COLORS") == "0":
return False
if os.environ.get("PYTHON_COLORS") == "1":
return True
if "NO_COLOR" in os.environ:
return False
if not COLORIZE:
return False
if not sys.flags.ignore_environment:
if "FORCE_COLOR" in os.environ:
return True
if os.environ.get("TERM") == "dumb":
return False

if not hasattr(sys.stderr, "fileno"):
return False

try:
return os.isatty(sys.stderr.fileno())
except io.UnsupportedOperation:
return sys.stderr.isatty()
38 changes: 16 additions & 22 deletions Lib/doctest.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,8 @@ def _test():
import unittest
from io import StringIO, IncrementalNewlineDecoder
from collections import namedtuple
from traceback import _ANSIColors, _can_colorize
import _colorize # Used in doctests
from _colorize import ANSIColors, can_colorize


class TestResults(namedtuple('TestResults', 'failed attempted')):
Expand Down Expand Up @@ -1180,8 +1181,8 @@ class DocTestRunner:
The `run` method is used to process a single DocTest case. It
returns a TestResults instance.
>>> save_colorize = traceback._COLORIZE
>>> traceback._COLORIZE = False
>>> save_colorize = _colorize.COLORIZE
>>> _colorize.COLORIZE = False
>>> tests = DocTestFinder().find(_TestClass)
>>> runner = DocTestRunner(verbose=False)
Expand Down Expand Up @@ -1234,7 +1235,7 @@ class DocTestRunner:
overriding the methods `report_start`, `report_success`,
`report_unexpected_exception`, and `report_failure`.
>>> traceback._COLORIZE = save_colorize
>>> _colorize.COLORIZE = save_colorize
"""
# This divider string is used to separate failure messages, and to
# separate sections of the summary.
Expand Down Expand Up @@ -1314,7 +1315,7 @@ def report_unexpected_exception(self, out, test, example, exc_info):

def _failure_header(self, test, example):
red, reset = (
(_ANSIColors.RED, _ANSIColors.RESET) if _can_colorize() else ("", "")
(ANSIColors.RED, ANSIColors.RESET) if can_colorize() else ("", "")
)
out = [f"{red}{self.DIVIDER}{reset}"]
if test.filename:
Expand Down Expand Up @@ -1556,8 +1557,8 @@ def out(s):
# Make sure sys.displayhook just prints the value to stdout
save_displayhook = sys.displayhook
sys.displayhook = sys.__displayhook__
saved_can_colorize = traceback._can_colorize
traceback._can_colorize = lambda: False
saved_can_colorize = _colorize.can_colorize
_colorize.can_colorize = lambda: False
color_variables = {"PYTHON_COLORS": None, "FORCE_COLOR": None}
for key in color_variables:
color_variables[key] = os.environ.pop(key, None)
Expand All @@ -1569,7 +1570,7 @@ def out(s):
sys.settrace(save_trace)
linecache.getlines = self.save_linecache_getlines
sys.displayhook = save_displayhook
traceback._can_colorize = saved_can_colorize
_colorize.can_colorize = saved_can_colorize
for key, value in color_variables.items():
if value is not None:
os.environ[key] = value
Expand Down Expand Up @@ -1609,20 +1610,13 @@ def summarize(self, verbose=None):
else:
failed.append((name, (failures, tries, skips)))

if _can_colorize():
bold_green = _ANSIColors.BOLD_GREEN
bold_red = _ANSIColors.BOLD_RED
green = _ANSIColors.GREEN
red = _ANSIColors.RED
reset = _ANSIColors.RESET
yellow = _ANSIColors.YELLOW
else:
bold_green = ""
bold_red = ""
green = ""
red = ""
reset = ""
yellow = ""
ansi = _colorize.get_colors()
bold_green = ansi.BOLD_GREEN
bold_red = ansi.BOLD_RED
green = ansi.GREEN
red = ansi.RED
reset = ansi.RESET
yellow = ansi.YELLOW

if verbose:
if notests:
Expand Down
9 changes: 5 additions & 4 deletions Lib/test/support/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2579,20 +2579,21 @@ def copy_python_src_ignore(path, names):
}
return ignored


def force_not_colorized(func):
"""Force the terminal not to be colorized."""
@functools.wraps(func)
def wrapper(*args, **kwargs):
import traceback
original_fn = traceback._can_colorize
import _colorize
original_fn = _colorize.can_colorize
variables = {"PYTHON_COLORS": None, "FORCE_COLOR": None}
try:
for key in variables:
variables[key] = os.environ.pop(key, None)
traceback._can_colorize = lambda: False
_colorize.can_colorize = lambda: False
return func(*args, **kwargs)
finally:
traceback._can_colorize = original_fn
_colorize.can_colorize = original_fn
for key, value in variables.items():
if value is not None:
os.environ[key] = value
Expand Down
59 changes: 59 additions & 0 deletions Lib/test/test__colorize.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
import contextlib
import sys
import unittest
import unittest.mock
import _colorize
from test.support import force_not_colorized

ORIGINAL_CAN_COLORIZE = _colorize.can_colorize


def setUpModule():
_colorize.can_colorize = lambda: False


def tearDownModule():
_colorize.can_colorize = ORIGINAL_CAN_COLORIZE


class TestColorizeFunction(unittest.TestCase):
@force_not_colorized
def test_colorized_detection_checks_for_environment_variables(self):
if sys.platform == "win32":
virtual_patching = unittest.mock.patch("nt._supports_virtual_terminal",
return_value=True)
else:
virtual_patching = contextlib.nullcontext()
with virtual_patching:

flags = unittest.mock.MagicMock(ignore_environment=False)
with (unittest.mock.patch("os.isatty") as isatty_mock,
unittest.mock.patch("sys.flags", flags),
unittest.mock.patch("_colorize.can_colorize", ORIGINAL_CAN_COLORIZE)):
isatty_mock.return_value = True
with unittest.mock.patch("os.environ", {'TERM': 'dumb'}):
self.assertEqual(_colorize.can_colorize(), False)
with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '1'}):
self.assertEqual(_colorize.can_colorize(), True)
with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '0'}):
self.assertEqual(_colorize.can_colorize(), False)
with unittest.mock.patch("os.environ", {'NO_COLOR': '1'}):
self.assertEqual(_colorize.can_colorize(), False)
with unittest.mock.patch("os.environ",
{'NO_COLOR': '1', "PYTHON_COLORS": '1'}):
self.assertEqual(_colorize.can_colorize(), True)
with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1'}):
self.assertEqual(_colorize.can_colorize(), True)
with unittest.mock.patch("os.environ",
{'FORCE_COLOR': '1', 'NO_COLOR': '1'}):
self.assertEqual(_colorize.can_colorize(), False)
with unittest.mock.patch("os.environ",
{'FORCE_COLOR': '1', "PYTHON_COLORS": '0'}):
self.assertEqual(_colorize.can_colorize(), False)
isatty_mock.return_value = False
with unittest.mock.patch("os.environ", {}):
self.assertEqual(_colorize.can_colorize(), False)


if __name__ == "__main__":
unittest.main()

0 comments on commit 1e0a640

Please sign in to comment.