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

Author: hugovk

Lib/_colorize.py

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

Lib/doctest.py

@@ -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')):
@@ -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)
@@ -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.
@@ -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:
@@ -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)
@@ -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
@@ -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:

Lib/test/support/__init__.py

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

Lib/test/test__colorize.py

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