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

1 files changed, 27 insertions, 69 deletions

diff --git a/Lib/traceback.py b/Lib/traceback.py
index 6ce745f..8403173 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
@@ -9,6 +7,8 @@ import sys
 import textwrap
 import warnings
 from contextlib import suppress
+import _colorize
+from _colorize import ANSIColors
 
 __all__ = ['extract_stack', 'extract_tb', 'format_exception',
            '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.
 #
 
-_COLORIZE = True
 
 def print_list(extracted_list, file=None):
     """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()
 
-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, /):
     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)
 
 
@@ -214,9 +182,9 @@ def _format_final_exc_line(etype, value, *, insert_final_newline=True, colorize=
     end_char = "\n" if insert_final_newline else ""
     if colorize:
         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:
-            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:
         if value is None or not valuestr:
             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}"
     return line
 
+
 def _safe_string(value, what, func=str):
     try:
         return func(value)
@@ -449,17 +418,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'
-    GREEN = "\x1b[32m"
-    BOLD_GREEN = "\x1b[1;32m"
-    GREY = '\x1b[90m'
-    RESET = '\x1b[0m'
-    YELLOW = "\x1b[33m"
-
 
 class StackSummary(list):
     """A list of FrameSummary objects, representing a stack of frames."""
@@ -564,15 +522,15 @@ class StackSummary(list):
             filename = "<stdin>"
         if colorize:
             row.append('  File {}"{}"{}, line {}{}{}, in {}{}{}\n'.format(
-                    _ANSIColors.MAGENTA,
+                    ANSIColors.MAGENTA,
                     filename,
-                    _ANSIColors.RESET,
-                    _ANSIColors.MAGENTA,
+                    ANSIColors.RESET,
+                    ANSIColors.MAGENTA,
                     frame_summary.lineno,
-                    _ANSIColors.RESET,
-                    _ANSIColors.MAGENTA,
+                    ANSIColors.RESET,
+                    ANSIColors.MAGENTA,
                     frame_summary.name,
-                    _ANSIColors.RESET,
+                    ANSIColors.RESET,
                     )
             )
         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]):
                             caret_group = list(group)
                             if color == "^":
-                                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_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)
                             elif color == "~":
-                                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_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)
                             else:
                                 colorized_line_parts.append("".join(char for char, _ 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 colorize:
                 yield '  File {}"{}"{}, line {}{}{}\n'.format(
-                    _ANSIColors.MAGENTA,
+                    ANSIColors.MAGENTA,
                     self.filename or "<string>",
-                    _ANSIColors.RESET,
-                    _ANSIColors.MAGENTA,
+                    ANSIColors.RESET,
+                    ANSIColors.MAGENTA,
                     self.lineno,
-                    _ANSIColors.RESET,
+                    ANSIColors.RESET,
                     )
             else:
                 yield '  File "{}", line {}\n'.format(
@@ -1352,11 +1310,11 @@ class TracebackException:
                         # colorize from colno to end_colno
                         ltext = (
                             ltext[:colno] +
-                            _ANSIColors.BOLD_RED + ltext[colno:end_colno] + _ANSIColors.RESET +
+                            ANSIColors.BOLD_RED + ltext[colno:end_colno] + ANSIColors.RESET +
                             ltext[end_colno:]
                         )
-                        start_color = _ANSIColors.BOLD_RED
-                        end_color = _ANSIColors.RESET
+                        start_color = ANSIColors.BOLD_RED
+                        end_color = ANSIColors.RESET
                     yield '    {}\n'.format(ltext)
                     yield '    {}{}{}{}\n'.format(
                         "".join(caretspace),
@@ -1369,12 +1327,12 @@ class TracebackException:
         msg = self.msg or "<no detail available>"
         if colorize:
             yield "{}{}{}: {}{}{}{}\n".format(
-                _ANSIColors.BOLD_MAGENTA,
+                ANSIColors.BOLD_MAGENTA,
                 stype,
-                _ANSIColors.RESET,
-                _ANSIColors.MAGENTA,
+                ANSIColors.RESET,
+                ANSIColors.MAGENTA,
                 msg,
-                _ANSIColors.RESET,
+                ANSIColors.RESET,
                 filename_suffix)
         else:
             yield "{}: {}{}\n".format(stype, msg, filename_suffix)