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,7 @@ def _test():
 import unittest
 from io import StringIO, IncrementalNewlineDecoder
 from collections import namedtuple
+import _colorize  # Used in doctests
 
 
 class TestResults(namedtuple('TestResults', 'failed attempted')):
@@ -1030,7 +1031,7 @@ def _find(self, tests, obj, name, module, source_lines, globs, seen):
         # Look for tests in a module's contained objects.
         if inspect.ismodule(obj) and self._recurse:
             for valname, val in obj.__dict__.items():
-                valname = '%s.%s' % (name, valname)
+                valname = f'{name}.{valname}'
 
                 # Recurse to functions & classes.
                 if ((self._is_routine(val) or inspect.isclass(val)) and
@@ -1051,7 +1052,7 @@ def _find(self, tests, obj, name, module, source_lines, globs, seen):
                                      "must be strings, functions, methods, "
                                      "classes, or modules: %r" %
                                      (type(val),))
-                valname = '%s.__test__.%s' % (name, valname)
+                valname = f'{name}.__test__.{valname}'
                 self._find(tests, val, valname, module, source_lines,
                            globs, seen)
 
@@ -1066,7 +1067,7 @@ def _find(self, tests, obj, name, module, source_lines, globs, seen):
                 if ((inspect.isroutine(val) or inspect.isclass(val) or
                       isinstance(val, property)) and
                       self._from_module(module, val)):
-                    valname = '%s.%s' % (name, valname)
+                    valname = f'{name}.{valname}'
                     self._find(tests, val, valname, module, source_lines,
                                globs, seen)
 
@@ -1172,6 +1173,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 +1226,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.
@@ -1309,7 +1315,7 @@ def _failure_header(self, test, example):
             out.append('File "%s", line %s, in %s' %
                        (test.filename, lineno, test.name))
         else:
-            out.append('Line %s, in %s' % (example.lineno+1, test.name))
+            out.append(f'Line {example.lineno+1}, in {test.name}')
         out.append('Failed example:')
         source = example.source
         out.append(_indent(source))
@@ -1566,10 +1572,12 @@ def summarize(self, verbose=None):
         summary is.  If the verbosity is not specified, then the
         DocTestRunner's verbosity is used.
         """
+        from _colorize import ANSIColors, can_colorize
+
         if verbose is None:
             verbose = self._verbose
 
-        notests, passed, failed = [], [], []
+        no_tests, passed, failed = [], [], []
         total_tries = total_failures = total_skips = 0
 
         for name, (failures, tries, skips) in self._stats.items():
@@ -1579,47 +1587,65 @@ def summarize(self, verbose=None):
             total_skips += skips
 
             if tries == 0:
-                notests.append(name)
+                no_tests.append(name)
             elif failures == 0:
                 passed.append((name, tries))
             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:")
-                notests.sort()
-                for name in notests:
+            if no_tests:
+                print(f"{_n_items(no_tests)} had no tests:")
+                no_tests.sort()
+                for name in no_tests:
                     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"{_n_items(failed)} had failures:")
+            print(f"{red}{_n_items(failed)} had failures:{reset}")
             for name, (failures, tries, skips) in sorted(failed):
-                print(f" {failures:3d} of {tries:3d} in {name}")
+                print(f"{red} {failures:3d} of {tries:3d} in {name}{reset}")
 
         if verbose:
             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 +1663,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 +1674,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
@@ -1784,7 +1810,7 @@ def output_difference(self, example, got, optionflags):
         # If we're not using diff, then simply list the expected
         # output followed by the actual output.
         if want and got:
-            return 'Expected:\n%sGot:\n%s' % (_indent(want), _indent(got))
+            return f'Expected:\n{_indent(want)}Got:\n{_indent(got)}'
         elif want:
             return 'Expected:\n%sGot nothing\n' % _indent(want)
         elif got:
@@ -2019,7 +2045,7 @@ class doctest.Tester, then merges the results into (or creates)
 
     # Check that we were actually given a module.
     if not inspect.ismodule(m):
-        raise TypeError("testmod: module required; %r" % (m,))
+        raise TypeError(f"testmod: module required; {m!r}")
 
     # If no name was given, then use the module's name.
     if name is None:
@@ -2393,7 +2419,7 @@ def __hash__(self):
 
     def __repr__(self):
         name = self._dt_test.name.split('.')
-        return "%s (%s)" % (name[-1], '.'.join(name[:-1]))
+        return f"{name[-1]} ({'.'.join(name[:-1])})"
 
     __str__ = object.__str__
 

Lib/test/test_doctest/test_doctest.py

@@ -16,6 +16,7 @@
 import tempfile
 import types
 import contextlib
+import _colorize  # used in doctests
 
 
 if not support.has_subprocess_support:
@@ -466,7 +467,7 @@ def basics(): r"""
     >>> tests = finder.find(sample_func)
 
     >>> print(tests)  # doctest: +ELLIPSIS
-    [<DocTest sample_func from test_doctest.py:33 (1 example)>]
+    [<DocTest sample_func from test_doctest.py:35 (1 example)>]
 
 The exact name depends on how test_doctest was invoked, so allow for
 leading path components.
@@ -2634,8 +2635,10 @@ def test_testfile(): r"""
 called with the name of a file, which is taken to be relative to the
 calling module.  The return value is (#failures, #tests).
 
-We don't want `-v` in sys.argv for these tests.
+We don't want color or `-v` in sys.argv for these tests.
 
+    >>> save_colorize = _colorize.COLORIZE
+    >>> _colorize.COLORIZE = False
     >>> save_argv = sys.argv
     >>> if '-v' in sys.argv:
     ...     sys.argv = [arg for arg in save_argv if arg != '-v']
@@ -2802,6 +2805,7 @@ def test_testfile(): r"""
     TestResults(failed=0, attempted=2)
     >>> doctest.master = None  # Reset master.
     >>> sys.argv = save_argv
+    >>> _colorize.COLORIZE = save_colorize
 """
 
 class TestImporter(importlib.abc.MetaPathFinder, importlib.abc.ResourceLoader):

Lib/test/test_doctest/test_doctest2.py

@@ -13,6 +13,8 @@
 
 import sys
 import unittest
+import _colorize
+
 if sys.flags.optimize >= 2:
     raise unittest.SkipTest("Cannot test docstrings with -O2")
 
@@ -108,14 +110,24 @@ def clsm(cls, val):
 
 
 class Test(unittest.TestCase):
+    def setUp(self):
+        super().setUp()
+        self.colorize = _colorize.COLORIZE
+        _colorize.COLORIZE = False
+
+    def tearDown(self):
+        super().tearDown()
+        _colorize.COLORIZE = self.colorize
+
     def test_testmod(self):
-        import doctest, sys
+        import doctest
+        import sys
         EXPECTED = 19
         f, t = doctest.testmod(sys.modules[__name__])
         if f:
-            self.fail("%d of %d doctests failed" % (f, t))
+            self.fail(f"{f} of {t} doctests failed")
         if t != EXPECTED:
-            self.fail("expected %d tests to run, not %d" % (EXPECTED, t))
+            self.fail(f"expected {EXPECTED} tests to run, not {t}")
 
 
 # Pollute the namespace with a bunch of imported functions and classes,