Use _colorize in traceback module

python · hugovk · Apr 24, 2024 · Apr 6, 2024 · Apr 6, 2024 · Apr 6, 2024
commit d27c0a853d44d768ef7e0f7fdcaacf0ab04035ae
diff --git 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}"<string>"{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
@@ -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,19 +8,20 @@
 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',
            'print_last', 'print_stack', 'print_tb', 'clear_frames',
            '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."""