diff --git a/Lib/test/test_traceback.py b/Lib/test/test_traceback.py index dd9b1850adf086..2f55d2392c484b 100644 --- a/Lib/test/test_traceback.py +++ b/Lib/test/test_traceback.py @@ -29,6 +29,8 @@ from functools import partial from pathlib import Path +import _colorize + MODULE_PREFIX = f'{__name__}.' if __name__ == '__main__' else '' test_code = namedtuple('code', ['co_filename', 'co_name']) @@ -45,12 +47,12 @@ class TracebackCases(unittest.TestCase): # formatting of SyntaxErrors works based on changes for 2.1. def setUp(self): super().setUp() - self.colorize = traceback._COLORIZE - traceback._COLORIZE = False + self.colorize = _colorize._COLORIZE + _colorize._COLORIZE = False def tearDown(self): super().tearDown() - traceback._COLORIZE = self.colorize + _colorize._COLORIZE = self.colorize def get_exception_format(self, func, exc): try: @@ -4291,9 +4293,9 @@ def bar(): e, capture_locals=True ) lines = "".join(exc.format(colorize=True)) - red = traceback._ANSIColors.RED - boldr = traceback._ANSIColors.BOLD_RED - reset = traceback._ANSIColors.RESET + red = _colorize._ANSIColors.RED + boldr = _colorize._ANSIColors.BOLD_RED + reset = _colorize._ANSIColors.RESET self.assertIn("y = " + red + "x['a']['b']" + reset + boldr + "['c']" + reset, lines) self.assertIn("return " + red + "foo" + reset + boldr + "(1,2,3,4)" + reset, lines) self.assertIn("return " + red + "baz" + reset + boldr + "(1," + reset, lines) @@ -4309,11 +4311,11 @@ def test_colorized_syntax_error(self): e, capture_locals=True ) actual = "".join(exc.format(colorize=True)) - red = traceback._ANSIColors.RED - magenta = traceback._ANSIColors.MAGENTA - boldm = traceback._ANSIColors.BOLD_MAGENTA - boldr = traceback._ANSIColors.BOLD_RED - reset = traceback._ANSIColors.RESET + red = _colorize._ANSIColors.RED + magenta = _colorize._ANSIColors.MAGENTA + boldm = _colorize._ANSIColors.BOLD_MAGENTA + boldr = _colorize._ANSIColors.BOLD_RED + reset = _colorize._ANSIColors.RESET expected = "".join([ f' File {magenta}""{reset}, line {magenta}1{reset}\n', f' a {boldr}${reset} b\n', @@ -4332,15 +4334,15 @@ def foo(): self.fail("No exception thrown.") except Exception as e: 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) actual = tbstderr.getvalue().splitlines() - red = traceback._ANSIColors.RED - boldr = traceback._ANSIColors.BOLD_RED - magenta = traceback._ANSIColors.MAGENTA - boldm = traceback._ANSIColors.BOLD_MAGENTA - reset = traceback._ANSIColors.RESET + 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):', f' File {magenta}"{__file__}"{reset}, ' @@ -4363,23 +4365,24 @@ def test_colorized_detection_checks_for_environment_variables(self): with unittest.mock.patch("os.isatty") as isatty_mock: isatty_mock.return_value = True with unittest.mock.patch("os.environ", {'TERM': 'dumb'}): - self.assertEqual(traceback._can_colorize(), False) + self.assertEqual(_colorize._can_colorize(), False) with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '1'}): - self.assertEqual(traceback._can_colorize(), True) + self.assertEqual(_colorize._can_colorize(), True) with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '0'}): - self.assertEqual(traceback._can_colorize(), False) + self.assertEqual(_colorize._can_colorize(), False) with unittest.mock.patch("os.environ", {'NO_COLOR': '1'}): - self.assertEqual(traceback._can_colorize(), False) + self.assertEqual(_colorize._can_colorize(), False) with unittest.mock.patch("os.environ", {'NO_COLOR': '1', "PYTHON_COLORS": '1'}): - self.assertEqual(traceback._can_colorize(), True) + self.assertEqual(_colorize._can_colorize(), True) with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1'}): - self.assertEqual(traceback._can_colorize(), True) + self.assertEqual(_colorize._can_colorize(), True) with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1', 'NO_COLOR': '1'}): - self.assertEqual(traceback._can_colorize(), False) + self.assertEqual(_colorize._can_colorize(), False) with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1', "PYTHON_COLORS": '0'}): - self.assertEqual(traceback._can_colorize(), False) + self.assertEqual(_colorize._can_colorize(), False) isatty_mock.return_value = False - self.assertEqual(traceback._can_colorize(), False) + self.assertEqual(_colorize._can_colorize(), False) + if __name__ == "__main__": unittest.main() diff --git a/Lib/traceback.py b/Lib/traceback.py index d27c7a726d2bb6..a294fe21574c5a 100644 --- a/Lib/traceback.py +++ b/Lib/traceback.py @@ -1,7 +1,5 @@ """Extract, format and print information about Python stack traces.""" -import os -import io import collections.abc import itertools import linecache @@ -10,6 +8,8 @@ import warnings from contextlib import suppress +from _colorize import _ANSIColors + __all__ = ['extract_stack', 'extract_tb', 'format_exception', 'format_exception_only', 'format_list', 'format_stack', 'format_tb', 'print_exc', 'format_exc', 'print_exception', @@ -17,12 +17,11 @@ 'FrameSummary', 'StackSummary', 'TracebackException', 'walk_stack', 'walk_tb'] + # # Formatting and printing lists of traceback lines. # -_COLORIZE = True - def print_list(extracted_list, file=None): """Print the list of tuples as returned by extract_tb() or extract_stack() as a formatted stack trace to the given file.""" @@ -133,33 +132,10 @@ def print_exception(exc, /, value=_sentinel, tb=_sentinel, limit=None, \ 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 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 "FORCE_COLOR" in os.environ: - return True - if os.environ.get("TERM") == "dumb": - return False - try: - return os.isatty(sys.stderr.fileno()) - except io.UnsupportedOperation: - return sys.stderr.isatty() def _print_exception_bltin(exc, /): + from _colorize import _can_colorize + file = sys.stderr if sys.stderr is not None else sys.__stderr__ colorize = _can_colorize() return print_exception(exc, limit=BUILTIN_EXCEPTION_LIMIT, file=file, colorize=colorize) @@ -443,13 +419,6 @@ def _get_code_position(code, instruction_index): _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' - GREY = '\x1b[90m' - RESET = '\x1b[0m' class StackSummary(list): """A list of FrameSummary objects, representing a stack of frames."""