Add colour to doctest output and create _colorize module

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

@@ -105,6 +105,8 @@ def _test():
 from io import StringIO, IncrementalNewlineDecoder
 from collections import namedtuple
 
+import _colorize  # Used in doctests
+
 
 class TestResults(namedtuple('TestResults', 'failed attempted')):
     def __new__(cls, failed, attempted, *, skipped=0):
@@ -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.
@@ -1566,10 +1573,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 +1588,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 +1664,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 +1675,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/test_doctest/test_doctest.py

@@ -17,6 +17,8 @@
 import types
 import contextlib
 
+import _colorize  # used in doctests
+
 
 if not support.has_subprocess_support:
     raise unittest.SkipTest("test_CLI requires subprocess support.")
@@ -466,7 +468,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 +2636,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 colour 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 +2806,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,9 @@
 
 import sys
 import unittest
+
+import _colorize
+
 if sys.flags.optimize >= 2:
     raise unittest.SkipTest("Cannot test docstrings with -O2")
 
@@ -108,6 +111,15 @@ 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
         EXPECTED = 19