gh-117225: Add color to doctest output

Lib/_colorize.py

@@ -0,0 +1,45 @@
+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"
+
+
+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()

Lib/doctest.py

@@ -104,6 +104,8 @@ def _test():
 import unittest
 from io import StringIO, IncrementalNewlineDecoder
 from collections import namedtuple
+import _colorize  # Used in doctests
+from _colorize import ANSIColors, can_colorize
 
 
 class TestResults(namedtuple('TestResults', 'failed attempted')):
@@ -1172,6 +1174,9 @@ class DocTestRunner:
     The `run` method is used to process a single DocTest case.  It
     returns a TestResults instance.
 
+        >>> save_colorize = _colorize.COLORIZE
+        >>> _colorize.COLORIZE = False
+
         >>> tests = DocTestFinder().find(_TestClass)
         >>> runner = DocTestRunner(verbose=False)
         >>> tests.sort(key = lambda test: test.name)
@@ -1222,6 +1227,8 @@ class DocTestRunner:
     can be also customized by subclassing DocTestRunner, and
     overriding the methods `report_start`, `report_success`,
     `report_unexpected_exception`, and `report_failure`.
+
+        >>> _colorize.COLORIZE = save_colorize
     """
     # This divider string is used to separate failure messages, and to
     # separate sections of the summary.
@@ -1300,7 +1307,10 @@ def report_unexpected_exception(self, out, test, example, exc_info):
             'Exception raised:\n' + _indent(_exception_traceback(exc_info)))
 
     def _failure_header(self, test, example):
-        out = [self.DIVIDER]
+        red, reset = (
+            (ANSIColors.RED, ANSIColors.RESET) if can_colorize() else ("", "")
+        )
+        out = [f"{red}{self.DIVIDER}{reset}"]
         if test.filename:
             if test.lineno is not None and example.lineno is not None:
                 lineno = test.lineno + example.lineno + 1
@@ -1585,6 +1595,21 @@ 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 = ""
+
         if verbose:
             if notests:
                 print(f"{_n_items(notests)} had no tests:")
@@ -1593,13 +1618,13 @@ def summarize(self, verbose=None):
                     print(f"    {name}")
 
             if passed:
-                print(f"{_n_items(passed)} passed all tests:")
+                print(f"{green}{_n_items(passed)} passed all tests:{reset}")
                 for name, count in sorted(passed):
                     s = "" if count == 1 else "s"
-                    print(f" {count:3d} test{s} in {name}")
+                    print(f" {green}{count:3d} test{s} in {name}{reset}")
 
         if failed:
-            print(self.DIVIDER)
+            print(f"{red}{self.DIVIDER}{reset}")
             print(f"{_n_items(failed)} had failures:")
             for name, (failures, tries, skips) in sorted(failed):
                 print(f" {failures:3d} of {tries:3d} in {name}")
@@ -1608,18 +1633,21 @@ def summarize(self, verbose=None):
             s = "" if total_tries == 1 else "s"
             print(f"{total_tries} test{s} in {_n_items(self._stats)}.")
 
-            and_f = f" and {total_failures} failed" if total_failures else ""
-            print(f"{total_tries - total_failures} passed{and_f}.")
+            and_f = (
+                f" and {red}{total_failures} failed{reset}"
+                if total_failures else ""
+            )
+            print(f"{green}{total_tries - total_failures} passed{reset}{and_f}.")
 
         if total_failures:
             s = "" if total_failures == 1 else "s"
-            msg = f"***Test Failed*** {total_failures} failure{s}"
+            msg = f"{bold_red}***Test Failed*** {total_failures} failure{s}{reset}"
             if total_skips:
                 s = "" if total_skips == 1 else "s"
-                msg = f"{msg} and {total_skips} skipped test{s}"
+                msg = f"{msg} and {yellow}{total_skips} skipped test{s}{reset}"
             print(f"{msg}.")
         elif verbose:
-            print("Test passed.")
+            print(f"{bold_green}Test passed.{reset}")
 
         return TestResults(total_failures, total_tries, skipped=total_skips)
 
@@ -1637,7 +1665,7 @@ def merge(self, other):
             d[name] = (failures, tries, skips)
 
 
-def _n_items(items: list) -> str:
+def _n_items(items: list | dict) -> str:
     """
     Helper to pluralise the number of items in a list.
     """
@@ -1648,7 +1676,7 @@ def _n_items(items: list) -> str:
 
 class OutputChecker:
     """
-    A class used to check the whether the actual output from a doctest
+    A class used to check whether the actual output from a doctest
     example matches the expected output.  `OutputChecker` defines two
     methods: `check_output`, which compares a given pair of outputs,
     and returns true if they match; and `output_difference`, which

Lib/test/support/__init__.py

@@ -26,7 +26,7 @@
     "Error", "TestFailed", "TestDidNotRun", "ResourceDenied",
     # io
     "record_original_stdout", "get_original_stdout", "captured_stdout",
-    "captured_stdin", "captured_stderr",
+    "captured_stdin", "captured_stderr", "captured_output",
     # unittest
     "is_resource_enabled", "requires", "requires_freebsd_version",
     "requires_gil_enabled", "requires_linux_version", "requires_mac_ver",

Lib/test/test__colorize.py

@@ -0,0 +1,134 @@
+import contextlib
+import sys
+import traceback
+import unittest
+import unittest.mock
+import _colorize
+from test.support import captured_output
+
+
+class TestColorizedTraceback(unittest.TestCase):
+    def test_colorized_traceback(self):
+        def foo(*args):
+            x = {'a':{'b': None}}
+            y = x['a']['b']['c']
+
+        def baz(*args):
+            return foo(1,2,3,4)
+
+        def bar():
+            return baz(1,
+                    2,3
+                    ,4)
+        try:
+            bar()
+        except Exception as e:
+            exc = traceback.TracebackException.from_exception(
+                e, capture_locals=True
+            )
+        lines = "".join(exc.format(colorize=True))
+        red = _colorize.ANSIColors.RED
+        boldr = _colorize.ANSIColors.BOLD_RED
+        reset = _colorize.ANSIColors.RESET
+
+        expected = [
+            "y = " + red + "x['a']['b']" + reset + boldr + "['c']" + reset,
+            "return " + red + "foo" + reset + boldr + "(1,2,3,4)" + reset,
+            "return " + red + "baz" + reset + boldr + "(1," + reset,
+            boldr + "2,3" + reset,
+            boldr + ",4)" + reset,
+            red + "bar" + reset + boldr + "()" + reset,
+        ]
+        for line in expected:
+            self.assertIn(line, lines)
+
+    def test_colorized_syntax_error(self):
+        try:
+            compile("a $ b", "<string>", "exec")
+        except SyntaxError as e:
+            exc = traceback.TracebackException.from_exception(
+                e, capture_locals=True
+            )
+        actual = "".join(exc.format(colorize=True))
+        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',
+                f'      {boldr}^{reset}\n',
+                f'{boldm}SyntaxError{reset}: {magenta}invalid syntax{reset}\n',
+            ]
+        )
+        self.assertIn(expected, actual)
+
+    def test_colorized_traceback_is_the_default(self):
+        def foo():
+            1/0
+
+        from _testcapi import exception_print
+
+        try:
+            foo()
+            self.fail("No exception thrown.")
+        except Exception as e:
+            with captured_output("stderr") as tbstderr:
+                with unittest.mock.patch('_colorize.can_colorize', return_value=True):
+                    exception_print(e)
+            actual = tbstderr.getvalue().splitlines()
+
+        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}, '
+            f'line {magenta}{lno_foo+6}{reset}, in {magenta}test_colorized_traceback_is_the_default{reset}',
+            f'    {red}foo{reset+boldr}(){reset}',
+            f'    {red}~~~{reset+boldr}^^{reset}',
+            f'  File {magenta}"{__file__}"{reset}, '
+            f'line {magenta}{lno_foo+1}{reset}, in {magenta}foo{reset}',
+            f'    {red}1{reset+boldr}/{reset+red}0{reset}',
+            f'    {red}~{reset+boldr}^{reset+red}~{reset}',
+            f'{boldm}ZeroDivisionError{reset}: {magenta}division by zero{reset}',
+        ]
+        self.assertEqual(actual, expected)
+
+    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()
+
+        env_vars_expected = [
+            ({'TERM': 'dumb'}, False),
+            ({'PYTHON_COLORS': '1'}, True),
+            ({'PYTHON_COLORS': '0'}, False),
+            ({'NO_COLOR': '1'}, False),
+            ({'NO_COLOR': '1', "PYTHON_COLORS": '1'}, True),
+            ({'FORCE_COLOR': '1'}, True),
+            ({'FORCE_COLOR': '1', 'NO_COLOR': '1'}, False),
+            ({'FORCE_COLOR': '1', "PYTHON_COLORS": '0'}, False),
+        ]
+
+        with virtual_patching:
+            with unittest.mock.patch("os.isatty") as isatty_mock:
+                isatty_mock.return_value = True
+
+                for env_vars, expected in env_vars_expected:
+                    with unittest.mock.patch("os.environ", env_vars):
+                        self.assertEqual(_colorize.can_colorize(), expected)
+
+                isatty_mock.return_value = False
+                self.assertEqual(_colorize.can_colorize(), False)
+
+
+if __name__ == "__main__":
+    unittest.main()