mirrors/cpython
mirrors/cpython
0
0
Fork 0
mirror of https://github.com/python/cpython.git synced 2024-11-24 00:38:00 +01:00
Code Issues Releases Activity

gh-117225: Move colorize functionality to own internal module (#118283)

Browse Source
This commit is contained in:
Hugo van Kemenade 2024-05-01 21:27:06 +03:00 committed by GitHub
parent 164e2c31c0
commit 3b3f8dea57
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 218 additions and 180 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

64
Lib/_colorize.py Normal file
View File

@ -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
Lib/doctest.py
View File

@ -104,7 +104,8 @@ import traceback
import unittest import unittest
from io import StringIO, IncrementalNewlineDecoder from io import StringIO, IncrementalNewlineDecoder
from collections import namedtuple 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')): class TestResults(namedtuple('TestResults', 'failed attempted')):
@ -1180,8 +1181,8 @@ class DocTestRunner:
The `run` method is used to process a single DocTest case. It The `run` method is used to process a single DocTest case. It
returns a TestResults instance. returns a TestResults instance.
>>> save_colorize = traceback._COLORIZE >>> save_colorize = _colorize.COLORIZE
>>> traceback._COLORIZE = False >>> _colorize.COLORIZE = False
>>> tests = DocTestFinder().find(_TestClass) >>> tests = DocTestFinder().find(_TestClass)
>>> runner = DocTestRunner(verbose=False) >>> runner = DocTestRunner(verbose=False)
@ -1234,7 +1235,7 @@ class DocTestRunner:
overriding the methods `report_start`, `report_success`, overriding the methods `report_start`, `report_success`,
`report_unexpected_exception`, and `report_failure`. `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 # This divider string is used to separate failure messages, and to
# separate sections of the summary. # separate sections of the summary.
@ -1314,7 +1315,7 @@ class DocTestRunner:
def _failure_header(self, test, example): def _failure_header(self, test, example):
red, reset = ( red, reset = (
(_ANSIColors.RED, _ANSIColors.RESET) if _can_colorize() else ("", "") (ANSIColors.RED, ANSIColors.RESET) if can_colorize() else ("", "")
) )
out = [f"{red}{self.DIVIDER}{reset}"] out = [f"{red}{self.DIVIDER}{reset}"]
if test.filename: if test.filename:
@ -1556,8 +1557,8 @@ class DocTestRunner:
# Make sure sys.displayhook just prints the value to stdout # Make sure sys.displayhook just prints the value to stdout
save_displayhook = sys.displayhook save_displayhook = sys.displayhook
sys.displayhook = sys.__displayhook__ sys.displayhook = sys.__displayhook__
saved_can_colorize = traceback._can_colorize saved_can_colorize = _colorize.can_colorize
traceback._can_colorize = lambda: False _colorize.can_colorize = lambda: False
color_variables = {"PYTHON_COLORS": None, "FORCE_COLOR": None} color_variables = {"PYTHON_COLORS": None, "FORCE_COLOR": None}
for key in color_variables: for key in color_variables:
color_variables[key] = os.environ.pop(key, None) color_variables[key] = os.environ.pop(key, None)
@ -1569,7 +1570,7 @@ class DocTestRunner:
sys.settrace(save_trace) sys.settrace(save_trace)
linecache.getlines = self.save_linecache_getlines linecache.getlines = self.save_linecache_getlines
sys.displayhook = save_displayhook sys.displayhook = save_displayhook
traceback._can_colorize = saved_can_colorize _colorize.can_colorize = saved_can_colorize
for key, value in color_variables.items(): for key, value in color_variables.items():
if value is not None: if value is not None:
os.environ[key] = value os.environ[key] = value
@ -1609,20 +1610,13 @@ class DocTestRunner:
else: else:
failed.append((name, (failures, tries, skips))) failed.append((name, (failures, tries, skips)))
if _can_colorize(): ansi = _colorize.get_colors()
bold_green = _ANSIColors.BOLD_GREEN bold_green = ansi.BOLD_GREEN
bold_red = _ANSIColors.BOLD_RED bold_red = ansi.BOLD_RED
green = _ANSIColors.GREEN green = ansi.GREEN
red = _ANSIColors.RED red = ansi.RED
reset = _ANSIColors.RESET reset = ansi.RESET
yellow = _ANSIColors.YELLOW yellow = ansi.YELLOW
else:
bold_green = ""
bold_red = ""
green = ""
red = ""
reset = ""
yellow = ""
if verbose: if verbose:
if notests: if notests:

9
Lib/test/support/__init__.py
View File

@ -2579,20 +2579,21 @@ def copy_python_src_ignore(path, names):
} }
return ignored return ignored
def force_not_colorized(func): def force_not_colorized(func):
"""Force the terminal not to be colorized.""" """Force the terminal not to be colorized."""
@functools.wraps(func) @functools.wraps(func)
def wrapper(*args, **kwargs): def wrapper(*args, **kwargs):
import traceback import _colorize
original_fn = traceback._can_colorize original_fn = _colorize.can_colorize
variables = {"PYTHON_COLORS": None, "FORCE_COLOR": None} variables = {"PYTHON_COLORS": None, "FORCE_COLOR": None}
try: try:
for key in variables: for key in variables:
variables[key] = os.environ.pop(key, None) variables[key] = os.environ.pop(key, None)
traceback._can_colorize = lambda: False _colorize.can_colorize = lambda: False
return func(*args, **kwargs) return func(*args, **kwargs)
finally: finally:
traceback._can_colorize = original_fn _colorize.can_colorize = original_fn
for key, value in variables.items(): for key, value in variables.items():
if value is not None: if value is not None:
os.environ[key] = value os.environ[key] = value

59
Lib/test/test__colorize.py Normal file
View File

@ -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()

56
Lib/test/test_doctest/test_doctest.py
View File

@ -16,7 +16,7 @@ import unittest
import tempfile import tempfile
import types import types
import contextlib import contextlib
import traceback import _colorize
def doctest_skip_if(condition): def doctest_skip_if(condition):
@ -893,8 +893,8 @@ Unit tests for the `DocTestRunner` class.
DocTestRunner is used to run DocTest test cases, and to accumulate DocTestRunner is used to run DocTest test cases, and to accumulate
statistics. Here's a simple DocTest case we can use: statistics. Here's a simple DocTest case we can use:
>>> save_colorize = traceback._COLORIZE >>> save_colorize = _colorize.COLORIZE
>>> traceback._COLORIZE = False >>> _colorize.COLORIZE = False
>>> def f(x): >>> def f(x):
... ''' ... '''
@ -951,7 +951,7 @@ the failure and proceeds to the next example:
ok ok
TestResults(failed=1, attempted=3) TestResults(failed=1, attempted=3)
>>> traceback._COLORIZE = save_colorize >>> _colorize.COLORIZE = save_colorize
""" """
def verbose_flag(): r""" def verbose_flag(): r"""
The `verbose` flag makes the test runner generate more detailed The `verbose` flag makes the test runner generate more detailed
@ -1027,8 +1027,8 @@ An expected exception is specified with a traceback message. The
lines between the first line and the type/value may be omitted or lines between the first line and the type/value may be omitted or
replaced with any other string: replaced with any other string:
>>> save_colorize = traceback._COLORIZE >>> save_colorize = _colorize.COLORIZE
>>> traceback._COLORIZE = False >>> _colorize.COLORIZE = False
>>> def f(x): >>> def f(x):
... ''' ... '''
@ -1261,7 +1261,7 @@ unexpected exception:
ZeroDivisionError: integer division or modulo by zero ZeroDivisionError: integer division or modulo by zero
TestResults(failed=1, attempted=1) TestResults(failed=1, attempted=1)
>>> traceback._COLORIZE = save_colorize >>> _colorize.COLORIZE = save_colorize
""" """
def displayhook(): r""" def displayhook(): r"""
Test that changing sys.displayhook doesn't matter for doctest. Test that changing sys.displayhook doesn't matter for doctest.
@ -1303,8 +1303,8 @@ together).
The DONT_ACCEPT_TRUE_FOR_1 flag disables matches between True/False The DONT_ACCEPT_TRUE_FOR_1 flag disables matches between True/False
and 1/0: and 1/0:
>>> save_colorize = traceback._COLORIZE >>> save_colorize = _colorize.COLORIZE
>>> traceback._COLORIZE = False >>> _colorize.COLORIZE = False
>>> def f(x): >>> def f(x):
... '>>> True\n1\n' ... '>>> True\n1\n'
@ -1725,7 +1725,7 @@ more than one flag value. Here we verify that's fixed:
Clean up. Clean up.
>>> del doctest.OPTIONFLAGS_BY_NAME[unlikely] >>> del doctest.OPTIONFLAGS_BY_NAME[unlikely]
>>> traceback._COLORIZE = save_colorize >>> _colorize.COLORIZE = save_colorize
""" """
@ -1736,8 +1736,8 @@ Option directives can be used to turn option flags on or off for a
single example. To turn an option on for an example, follow that single example. To turn an option on for an example, follow that
example with a comment of the form ``# doctest: +OPTION``: example with a comment of the form ``# doctest: +OPTION``:
>>> save_colorize = traceback._COLORIZE >>> save_colorize = _colorize.COLORIZE
>>> traceback._COLORIZE = False >>> _colorize.COLORIZE = False
>>> def f(x): r''' >>> def f(x): r'''
... >>> print(list(range(10))) # should fail: no ellipsis ... >>> print(list(range(10))) # should fail: no ellipsis
@ -1947,7 +1947,7 @@ source:
Traceback (most recent call last): Traceback (most recent call last):
ValueError: line 0 of the doctest for s has an option directive on a line with no example: '# doctest: +ELLIPSIS' ValueError: line 0 of the doctest for s has an option directive on a line with no example: '# doctest: +ELLIPSIS'
>>> traceback._COLORIZE = save_colorize >>> _colorize.COLORIZE = save_colorize
""" """
def test_testsource(): r""" def test_testsource(): r"""
@ -2031,8 +2031,8 @@ if not hasattr(sys, 'gettrace') or not sys.gettrace():
with a version that restores stdout. This is necessary for you to with a version that restores stdout. This is necessary for you to
see debugger output. see debugger output.
>>> save_colorize = traceback._COLORIZE >>> save_colorize = _colorize.COLORIZE
>>> traceback._COLORIZE = False >>> _colorize.COLORIZE = False
>>> doc = ''' >>> doc = '''
... >>> x = 42 ... >>> x = 42
@ -2157,7 +2157,7 @@ if not hasattr(sys, 'gettrace') or not sys.gettrace():
9 9
TestResults(failed=1, attempted=3) TestResults(failed=1, attempted=3)
>>> traceback._COLORIZE = save_colorize >>> _colorize.COLORIZE = save_colorize
""" """
def test_pdb_set_trace_nested(): def test_pdb_set_trace_nested():
@ -2694,8 +2694,8 @@ calling module. The return value is (#failures, #tests).
We don't want color or `-v` in sys.argv for these tests. We don't want color or `-v` in sys.argv for these tests.
>>> save_colorize = traceback._COLORIZE >>> save_colorize = _colorize.COLORIZE
>>> traceback._COLORIZE = False >>> _colorize.COLORIZE = False
>>> save_argv = sys.argv >>> save_argv = sys.argv
>>> if '-v' in sys.argv: >>> if '-v' in sys.argv:
@ -2863,7 +2863,7 @@ Test the verbose output:
TestResults(failed=0, attempted=2) TestResults(failed=0, attempted=2)
>>> doctest.master = None # Reset master. >>> doctest.master = None # Reset master.
>>> sys.argv = save_argv >>> sys.argv = save_argv
>>> traceback._COLORIZE = save_colorize >>> _colorize.COLORIZE = save_colorize
""" """
class TestImporter(importlib.abc.MetaPathFinder, importlib.abc.ResourceLoader): class TestImporter(importlib.abc.MetaPathFinder, importlib.abc.ResourceLoader):
@ -3001,8 +3001,8 @@ if supports_unicode:
def test_unicode(): """ def test_unicode(): """
Check doctest with a non-ascii filename: Check doctest with a non-ascii filename:
>>> save_colorize = traceback._COLORIZE >>> save_colorize = _colorize.COLORIZE
>>> traceback._COLORIZE = False >>> _colorize.COLORIZE = False
>>> doc = ''' >>> doc = '''
... >>> raise Exception('clé') ... >>> raise Exception('clé')
@ -3030,7 +3030,7 @@ Check doctest with a non-ascii filename:
Exception: clé Exception: clé
TestResults(failed=1, attempted=1) TestResults(failed=1, attempted=1)
>>> traceback._COLORIZE = save_colorize >>> _colorize.COLORIZE = save_colorize
""" """
@ -3325,8 +3325,8 @@ def test_run_doctestsuite_multiple_times():
def test_exception_with_note(note): def test_exception_with_note(note):
""" """
>>> save_colorize = traceback._COLORIZE >>> save_colorize = _colorize.COLORIZE
>>> traceback._COLORIZE = False >>> _colorize.COLORIZE = False
>>> test_exception_with_note('Note') >>> test_exception_with_note('Note')
Traceback (most recent call last): Traceback (most recent call last):
@ -3378,7 +3378,7 @@ def test_exception_with_note(note):
note note
TestResults(failed=1, attempted=...) TestResults(failed=1, attempted=...)
>>> traceback._COLORIZE = save_colorize >>> _colorize.COLORIZE = save_colorize
""" """
exc = ValueError('Text') exc = ValueError('Text')
exc.add_note(note) exc.add_note(note)
@ -3459,8 +3459,8 @@ def test_syntax_error_subclass_from_stdlib():
def test_syntax_error_with_incorrect_expected_note(): def test_syntax_error_with_incorrect_expected_note():
""" """
>>> save_colorize = traceback._COLORIZE >>> save_colorize = _colorize.COLORIZE
>>> traceback._COLORIZE = False >>> _colorize.COLORIZE = False
>>> def f(x): >>> def f(x):
... r''' ... r'''
@ -3491,7 +3491,7 @@ def test_syntax_error_with_incorrect_expected_note():
note2 note2
TestResults(failed=1, attempted=...) TestResults(failed=1, attempted=...)
>>> traceback._COLORIZE = save_colorize >>> _colorize.COLORIZE = save_colorize
""" """

75
Lib/test/test_traceback.py
View File

@ -26,9 +26,9 @@ from test.support import force_not_colorized
import json import json
import textwrap import textwrap
import traceback import traceback
import contextlib
from functools import partial from functools import partial
from pathlib import Path from pathlib import Path
import _colorize
MODULE_PREFIX = f'{__name__}.' if __name__ == '__main__' else '' MODULE_PREFIX = f'{__name__}.' if __name__ == '__main__' else ''
@ -40,25 +40,18 @@ test_tb = namedtuple('tb', ['tb_frame', 'tb_lineno', 'tb_next', 'tb_lasti'])
LEVENSHTEIN_DATA_FILE = Path(__file__).parent / 'levenshtein_examples.json' LEVENSHTEIN_DATA_FILE = Path(__file__).parent / 'levenshtein_examples.json'
ORIGINAL_CAN_COLORIZE = traceback._can_colorize
def setUpModule():
traceback._can_colorize = lambda: False
def tearDownModule():
traceback._can_colorize = ORIGINAL_CAN_COLORIZE
class TracebackCases(unittest.TestCase): class TracebackCases(unittest.TestCase):
# For now, a very minimal set of tests. I want to be sure that # For now, a very minimal set of tests. I want to be sure that
# formatting of SyntaxErrors works based on changes for 2.1. # formatting of SyntaxErrors works based on changes for 2.1.
def setUp(self): def setUp(self):
super().setUp() super().setUp()
self.colorize = traceback._COLORIZE self.colorize = _colorize.COLORIZE
traceback._COLORIZE = False _colorize.COLORIZE = False
def tearDown(self): def tearDown(self):
super().tearDown() super().tearDown()
traceback._COLORIZE = self.colorize _colorize.COLORIZE = self.colorize
def get_exception_format(self, func, exc): def get_exception_format(self, func, exc):
try: try:
@ -4478,9 +4471,9 @@ class TestColorizedTraceback(unittest.TestCase):
e, capture_locals=True e, capture_locals=True
) )
lines = "".join(exc.format(colorize=True)) lines = "".join(exc.format(colorize=True))
red = traceback._ANSIColors.RED red = _colorize.ANSIColors.RED
boldr = traceback._ANSIColors.BOLD_RED boldr = _colorize.ANSIColors.BOLD_RED
reset = traceback._ANSIColors.RESET reset = _colorize.ANSIColors.RESET
self.assertIn("y = " + red + "x['a']['b']" + reset + boldr + "['c']" + reset, lines) 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 " + 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) self.assertIn("return (lambda *args: " + red + "foo" + reset + boldr + "(*args)" + reset + ")(1,2,3,4)", lines)
@ -4496,11 +4489,11 @@ class TestColorizedTraceback(unittest.TestCase):
e, capture_locals=True e, capture_locals=True
) )
actual = "".join(exc.format(colorize=True)) actual = "".join(exc.format(colorize=True))
red = traceback._ANSIColors.RED red = _colorize.ANSIColors.RED
magenta = traceback._ANSIColors.MAGENTA magenta = _colorize.ANSIColors.MAGENTA
boldm = traceback._ANSIColors.BOLD_MAGENTA boldm = _colorize.ANSIColors.BOLD_MAGENTA
boldr = traceback._ANSIColors.BOLD_RED boldr = _colorize.ANSIColors.BOLD_RED
reset = traceback._ANSIColors.RESET reset = _colorize.ANSIColors.RESET
expected = "".join([ expected = "".join([
f' File {magenta}"<string>"{reset}, line {magenta}1{reset}\n', f' File {magenta}"<string>"{reset}, line {magenta}1{reset}\n',
f' a {boldr}${reset} b\n', f' a {boldr}${reset} b\n',
@ -4519,15 +4512,15 @@ class TestColorizedTraceback(unittest.TestCase):
self.fail("No exception thrown.") self.fail("No exception thrown.")
except Exception as e: except Exception as e:
with captured_output("stderr") as tbstderr: with captured_output("stderr") as tbstderr:
with unittest.mock.patch('traceback._can_colorize', return_value=True): with unittest.mock.patch('_colorize.can_colorize', return_value=True):
exception_print(e) exception_print(e)
actual = tbstderr.getvalue().splitlines() actual = tbstderr.getvalue().splitlines()
red = traceback._ANSIColors.RED red = _colorize.ANSIColors.RED
boldr = traceback._ANSIColors.BOLD_RED boldr = _colorize.ANSIColors.BOLD_RED
magenta = traceback._ANSIColors.MAGENTA magenta = _colorize.ANSIColors.MAGENTA
boldm = traceback._ANSIColors.BOLD_MAGENTA boldm = _colorize.ANSIColors.BOLD_MAGENTA
reset = traceback._ANSIColors.RESET reset = _colorize.ANSIColors.RESET
lno_foo = foo.__code__.co_firstlineno lno_foo = foo.__code__.co_firstlineno
expected = ['Traceback (most recent call last):', expected = ['Traceback (most recent call last):',
f' File {magenta}"{__file__}"{reset}, ' f' File {magenta}"{__file__}"{reset}, '
@ -4541,38 +4534,6 @@ class TestColorizedTraceback(unittest.TestCase):
f'{boldm}ZeroDivisionError{reset}: {magenta}division by zero{reset}'] f'{boldm}ZeroDivisionError{reset}: {magenta}division by zero{reset}']
self.assertEqual(actual, expected) self.assertEqual(actual, expected)
@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("traceback._can_colorize", ORIGINAL_CAN_COLORIZE)):
isatty_mock.return_value = True
with unittest.mock.patch("os.environ", {'TERM': 'dumb'}):
self.assertEqual(traceback._can_colorize(), False)
with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '1'}):
self.assertEqual(traceback._can_colorize(), True)
with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '0'}):
self.assertEqual(traceback._can_colorize(), False)
with unittest.mock.patch("os.environ", {'NO_COLOR': '1'}):
self.assertEqual(traceback._can_colorize(), False)
with unittest.mock.patch("os.environ", {'NO_COLOR': '1', "PYTHON_COLORS": '1'}):
self.assertEqual(traceback._can_colorize(), True)
with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1'}):
self.assertEqual(traceback._can_colorize(), True)
with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1', 'NO_COLOR': '1'}):
self.assertEqual(traceback._can_colorize(), False)
with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1', "PYTHON_COLORS": '0'}):
self.assertEqual(traceback._can_colorize(), False)
isatty_mock.return_value = False
with unittest.mock.patch("os.environ", {}):
self.assertEqual(traceback._can_colorize(), False)
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()

96
Lib/traceback.py
View File

@ -1,7 +1,5 @@
"""Extract, format and print information about Python stack traces.""" """Extract, format and print information about Python stack traces."""
import os
import io
import collections.abc import collections.abc
import itertools import itertools
import linecache import linecache
@ -9,6 +7,8 @@ import sys
import textwrap import textwrap
import warnings import warnings
from contextlib import suppress from contextlib import suppress
import _colorize
from _colorize import ANSIColors
__all__ = ['extract_stack', 'extract_tb', 'format_exception', __all__ = ['extract_stack', 'extract_tb', 'format_exception',
'format_exception_only', 'format_list', 'format_stack', 'format_exception_only', 'format_list', 'format_stack',
@ -21,7 +21,6 @@ __all__ = ['extract_stack', 'extract_tb', 'format_exception',
# Formatting and printing lists of traceback lines. # Formatting and printing lists of traceback lines.
# #
_COLORIZE = True
def print_list(extracted_list, file=None): def print_list(extracted_list, file=None):
"""Print the list of tuples as returned by extract_tb() or """Print the list of tuples as returned by extract_tb() or
@ -133,41 +132,10 @@ def print_exception(exc, /, value=_sentinel, tb=_sentinel, limit=None, \
BUILTIN_EXCEPTION_LIMIT = object() BUILTIN_EXCEPTION_LIMIT = object()
def _can_colorize():
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()
def _print_exception_bltin(exc, /): def _print_exception_bltin(exc, /):
file = sys.stderr if sys.stderr is not None else sys.__stderr__ file = sys.stderr if sys.stderr is not None else sys.__stderr__
colorize = _can_colorize() colorize = _colorize.can_colorize()
return print_exception(exc, limit=BUILTIN_EXCEPTION_LIMIT, file=file, colorize=colorize) return print_exception(exc, limit=BUILTIN_EXCEPTION_LIMIT, file=file, colorize=colorize)
@ -214,9 +182,9 @@ def _format_final_exc_line(etype, value, *, insert_final_newline=True, colorize=
end_char = "\n" if insert_final_newline else "" end_char = "\n" if insert_final_newline else ""
if colorize: if colorize:
if value is None or not valuestr: if value is None or not valuestr:
line = f"{_ANSIColors.BOLD_MAGENTA}{etype}{_ANSIColors.RESET}{end_char}" line = f"{ANSIColors.BOLD_MAGENTA}{etype}{ANSIColors.RESET}{end_char}"
else: else:
line = f"{_ANSIColors.BOLD_MAGENTA}{etype}{_ANSIColors.RESET}: {_ANSIColors.MAGENTA}{valuestr}{_ANSIColors.RESET}{end_char}" line = f"{ANSIColors.BOLD_MAGENTA}{etype}{ANSIColors.RESET}: {ANSIColors.MAGENTA}{valuestr}{ANSIColors.RESET}{end_char}"
else: else:
if value is None or not valuestr: if value is None or not valuestr:
line = f"{etype}{end_char}" line = f"{etype}{end_char}"
@ -224,6 +192,7 @@ def _format_final_exc_line(etype, value, *, insert_final_newline=True, colorize=
line = f"{etype}: {valuestr}{end_char}" line = f"{etype}: {valuestr}{end_char}"
return line return line
def _safe_string(value, what, func=str): def _safe_string(value, what, func=str):
try: try:
return func(value) return func(value)
@ -449,17 +418,6 @@ def _get_code_position(code, instruction_index):
_RECURSIVE_CUTOFF = 3 # Also hardcoded in traceback.c. _RECURSIVE_CUTOFF = 3 # Also hardcoded in traceback.c.
class _ANSIColors:
RED = '\x1b[31m'
BOLD_RED = '\x1b[1;31m'
MAGENTA = '\x1b[35m'
BOLD_MAGENTA = '\x1b[1;35m'
GREEN = "\x1b[32m"
BOLD_GREEN = "\x1b[1;32m"
GREY = '\x1b[90m'
RESET = '\x1b[0m'
YELLOW = "\x1b[33m"
class StackSummary(list): class StackSummary(list):
"""A list of FrameSummary objects, representing a stack of frames.""" """A list of FrameSummary objects, representing a stack of frames."""
@ -564,15 +522,15 @@ class StackSummary(list):
filename = "<stdin>" filename = "<stdin>"
if colorize: if colorize:
row.append(' File {}"{}"{}, line {}{}{}, in {}{}{}\n'.format( row.append(' File {}"{}"{}, line {}{}{}, in {}{}{}\n'.format(
_ANSIColors.MAGENTA, ANSIColors.MAGENTA,
filename, filename,
_ANSIColors.RESET, ANSIColors.RESET,
_ANSIColors.MAGENTA, ANSIColors.MAGENTA,
frame_summary.lineno, frame_summary.lineno,
_ANSIColors.RESET, ANSIColors.RESET,
_ANSIColors.MAGENTA, ANSIColors.MAGENTA,
frame_summary.name, frame_summary.name,
_ANSIColors.RESET, ANSIColors.RESET,
) )
) )
else: else:
@ -696,11 +654,11 @@ class StackSummary(list):
for color, group in itertools.groupby(itertools.zip_longest(line, carets, fillvalue=""), key=lambda x: x[1]): for color, group in itertools.groupby(itertools.zip_longest(line, carets, fillvalue=""), key=lambda x: x[1]):
caret_group = list(group) caret_group = list(group)
if color == "^": if color == "^":
colorized_line_parts.append(_ANSIColors.BOLD_RED + "".join(char for char, _ in caret_group) + _ANSIColors.RESET) colorized_line_parts.append(ANSIColors.BOLD_RED + "".join(char for char, _ in caret_group) + ANSIColors.RESET)
colorized_carets_parts.append(_ANSIColors.BOLD_RED + "".join(caret for _, caret in caret_group) + _ANSIColors.RESET) colorized_carets_parts.append(ANSIColors.BOLD_RED + "".join(caret for _, caret in caret_group) + ANSIColors.RESET)
elif color == "~": elif color == "~":
colorized_line_parts.append(_ANSIColors.RED + "".join(char for char, _ in caret_group) + _ANSIColors.RESET) colorized_line_parts.append(ANSIColors.RED + "".join(char for char, _ in caret_group) + ANSIColors.RESET)
colorized_carets_parts.append(_ANSIColors.RED + "".join(caret for _, caret in caret_group) + _ANSIColors.RESET) colorized_carets_parts.append(ANSIColors.RED + "".join(caret for _, caret in caret_group) + ANSIColors.RESET)
else: else:
colorized_line_parts.append("".join(char for char, _ in caret_group)) colorized_line_parts.append("".join(char for char, _ in caret_group))
colorized_carets_parts.append("".join(caret for _, caret in caret_group)) colorized_carets_parts.append("".join(caret for _, caret in caret_group))
@ -1307,12 +1265,12 @@ class TracebackException:
if self.lineno is not None: if self.lineno is not None:
if colorize: if colorize:
yield ' File {}"{}"{}, line {}{}{}\n'.format( yield ' File {}"{}"{}, line {}{}{}\n'.format(
_ANSIColors.MAGENTA, ANSIColors.MAGENTA,
self.filename or "<string>", self.filename or "<string>",
_ANSIColors.RESET, ANSIColors.RESET,
_ANSIColors.MAGENTA, ANSIColors.MAGENTA,
self.lineno, self.lineno,
_ANSIColors.RESET, ANSIColors.RESET,
) )
else: else:
yield ' File "{}", line {}\n'.format( yield ' File "{}", line {}\n'.format(
@ -1352,11 +1310,11 @@ class TracebackException:
# colorize from colno to end_colno # colorize from colno to end_colno
ltext = ( ltext = (
ltext[:colno] + ltext[:colno] +
_ANSIColors.BOLD_RED + ltext[colno:end_colno] + _ANSIColors.RESET + ANSIColors.BOLD_RED + ltext[colno:end_colno] + ANSIColors.RESET +
ltext[end_colno:] ltext[end_colno:]
) )
start_color = _ANSIColors.BOLD_RED start_color = ANSIColors.BOLD_RED
end_color = _ANSIColors.RESET end_color = ANSIColors.RESET
yield ' {}\n'.format(ltext) yield ' {}\n'.format(ltext)
yield ' {}{}{}{}\n'.format( yield ' {}{}{}{}\n'.format(
"".join(caretspace), "".join(caretspace),
@ -1369,12 +1327,12 @@ class TracebackException:
msg = self.msg or "<no detail available>" msg = self.msg or "<no detail available>"
if colorize: if colorize:
yield "{}{}{}: {}{}{}{}\n".format( yield "{}{}{}: {}{}{}{}\n".format(
_ANSIColors.BOLD_MAGENTA, ANSIColors.BOLD_MAGENTA,
stype, stype,
_ANSIColors.RESET, ANSIColors.RESET,
_ANSIColors.MAGENTA, ANSIColors.MAGENTA,
msg, msg,
_ANSIColors.RESET, ANSIColors.RESET,
filename_suffix) filename_suffix)
else: else:
yield "{}: {}{}\n".format(stype, msg, filename_suffix) yield "{}: {}{}\n".format(stype, msg, filename_suffix)

1
Python/stdlib_module_names.h generated
View File

@ -20,6 +20,7 @@ static const char* _Py_stdlib_module_names[] = {
"_codecs_tw", "_codecs_tw",
"_collections", "_collections",
"_collections_abc", "_collections_abc",
"_colorize",
"_compat_pickle", "_compat_pickle",
"_compression", "_compression",
"_contextvars", "_contextvars",