gh-117225: Add color to doctest output (#117583)

hugovk · AlexWaygood · web-flow · commit 975081b11e05 · 2024-04-24T14:27:40.000+03:00
Co-authored-by: Alex Waygood &lt;Alex.Waygood@Gmail.com&gt;
diff --git a/Lib/doctest.py b/Lib/doctest.py
@@ -104,6 +104,7 @@ def _test():
 import unittest
 from io import StringIO, IncrementalNewlineDecoder
 from collections import namedtuple
+from traceback import _ANSIColors, _can_colorize
 
 
 class TestResults(namedtuple('TestResults', 'failed attempted')):
@@ -1179,6 +1180,9 @@ 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
+
         >>> tests = DocTestFinder().find(_TestClass)
         >>> runner = DocTestRunner(verbose=False)
         >>> tests.sort(key = lambda test: test.name)
@@ -1229,6 +1233,8 @@ class DocTestRunner:
     can be also customized by subclassing DocTestRunner, and
     overriding the methods `report_start`, `report_success`,
     `report_unexpected_exception`, and `report_failure`.
+
+        >>> traceback._COLORIZE = save_colorize
     """
     # This divider string is used to separate failure messages, and to
     # separate sections of the summary.
@@ -1307,7 +1313,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
@@ -1592,6 +1601,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:")
@@ -1600,13 +1624,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}")
@@ -1615,18 +1639,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)
 
@@ -1644,7 +1671,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.
     """
@@ -1655,7 +1682,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
diff --git a/Lib/test/support/__init__.py b/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",
diff --git a/Lib/test/test_doctest/test_doctest.py b/Lib/test/test_doctest/test_doctest.py
@@ -16,6 +16,7 @@
 import tempfile
 import types
 import contextlib
+import traceback
 
 
 def doctest_skip_if(condition):
@@ -470,7 +471,7 @@ def basics(): r"""
     >>> tests = finder.find(sample_func)
 
     >>> print(tests)  # doctest: +ELLIPSIS
-    [<DocTest sample_func from test_doctest.py:37 (1 example)>]
+    [<DocTest sample_func from test_doctest.py:38 (1 example)>]
 
 The exact name depends on how test_doctest was invoked, so allow for
 leading path components.
@@ -892,6 +893,9 @@ def basics(): r"""
 DocTestRunner is used to run DocTest test cases, and to accumulate
 statistics.  Here's a simple DocTest case we can use:
 
+    >>> save_colorize = traceback._COLORIZE
+    >>> traceback._COLORIZE = False
+
     >>> def f(x):
     ...     '''
     ...     >>> x = 12
@@ -946,6 +950,8 @@ def basics(): r"""
         6
     ok
     TestResults(failed=1, attempted=3)
+
+    >>> traceback._COLORIZE = save_colorize
 """
     def verbose_flag(): r"""
 The `verbose` flag makes the test runner generate more detailed
@@ -1021,6 +1027,9 @@ def exceptions(): r"""
 lines between the first line and the type/value may be omitted or
 replaced with any other string:
 
+    >>> save_colorize = traceback._COLORIZE
+    >>> traceback._COLORIZE = False
+
     >>> def f(x):
     ...     '''
     ...     >>> x = 12
@@ -1251,6 +1260,8 @@ def exceptions(): r"""
         ...
         ZeroDivisionError: integer division or modulo by zero
     TestResults(failed=1, attempted=1)
+
+    >>> traceback._COLORIZE = save_colorize
 """
     def displayhook(): r"""
 Test that changing sys.displayhook doesn't matter for doctest.
@@ -1292,6 +1303,9 @@ def optionflags(): r"""
 The DONT_ACCEPT_TRUE_FOR_1 flag disables matches between True/False
 and 1/0:
 
+    >>> save_colorize = traceback._COLORIZE
+    >>> traceback._COLORIZE = False
+
     >>> def f(x):
     ...     '>>> True\n1\n'
 
@@ -1711,6 +1725,7 @@ def optionflags(): r"""
 
 Clean up.
     >>> del doctest.OPTIONFLAGS_BY_NAME[unlikely]
+    >>> traceback._COLORIZE = save_colorize
 
     """
 
@@ -1721,6 +1736,9 @@ def option_directives(): r"""
 single example.  To turn an option on for an example, follow that
 example with a comment of the form ``# doctest: +OPTION``:
 
+    >>> save_colorize = traceback._COLORIZE
+    >>> traceback._COLORIZE = False
+
     >>> def f(x): r'''
     ...     >>> print(list(range(10)))      # should fail: no ellipsis
     ...     [0, 1, ..., 9]
@@ -1928,6 +1946,8 @@ def option_directives(): r"""
     >>> test = doctest.DocTestParser().get_doctest(s, {}, 's', 's.py', 0)
     Traceback (most recent call last):
     ValueError: line 0 of the doctest for s has an option directive on a line with no example: '# doctest: +ELLIPSIS'
+
+    >>> traceback._COLORIZE = save_colorize
 """
 
 def test_testsource(): r"""
@@ -2011,6 +2031,9 @@ def test_pdb_set_trace():
         with a version that restores stdout.  This is necessary for you to
         see debugger output.
 
+          >>> save_colorize = traceback._COLORIZE
+          >>> traceback._COLORIZE = False
+
           >>> doc = '''
           ... >>> x = 42
           ... >>> raise Exception('clé')
@@ -2065,7 +2088,7 @@ def test_pdb_set_trace():
           ... finally:
           ...     sys.stdin = real_stdin
           --Return--
-          > <doctest test.test_doctest.test_doctest.test_pdb_set_trace[7]>(3)calls_set_trace()->None
+          > <doctest test.test_doctest.test_doctest.test_pdb_set_trace[9]>(3)calls_set_trace()->None
           -> import pdb; pdb.set_trace()
           (Pdb) print(y)
           2
@@ -2133,6 +2156,8 @@ def test_pdb_set_trace():
           Got:
               9
           TestResults(failed=1, attempted=3)
+
+          >>> traceback._COLORIZE = save_colorize
           """
 
     def test_pdb_set_trace_nested():
@@ -2667,7 +2692,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 = traceback._COLORIZE
+    >>> traceback._COLORIZE = False
 
     >>> save_argv = sys.argv
     >>> if '-v' in sys.argv:
@@ -2835,6 +2863,7 @@ def test_testfile(): r"""
     TestResults(failed=0, attempted=2)
     >>> doctest.master = None  # Reset master.
     >>> sys.argv = save_argv
+    >>> traceback._COLORIZE = save_colorize
 """
 
 class TestImporter(importlib.abc.MetaPathFinder, importlib.abc.ResourceLoader):
@@ -2972,6 +3001,9 @@ def test_testmod(): r"""
     def test_unicode(): """
 Check doctest with a non-ascii filename:
 
+    >>> save_colorize = traceback._COLORIZE
+    >>> traceback._COLORIZE = False
+
     >>> doc = '''
     ... >>> raise Exception('clé')
     ... '''
@@ -2997,8 +3029,11 @@ def test_unicode(): """
             raise Exception('clé')
         Exception: clé
     TestResults(failed=1, attempted=1)
+
+    >>> traceback._COLORIZE = save_colorize
     """
 
+
 @doctest_skip_if(not support.has_subprocess_support)
 def test_CLI(): r"""
 The doctest module can be used to run doctests against an arbitrary file.
@@ -3290,6 +3325,9 @@ def test_run_doctestsuite_multiple_times():
 
 def test_exception_with_note(note):
     """
+    >>> save_colorize = traceback._COLORIZE
+    >>> traceback._COLORIZE = False
+
     >>> test_exception_with_note('Note')
     Traceback (most recent call last):
       ...
@@ -3339,6 +3377,8 @@ def test_exception_with_note(note):
         ValueError: message
         note
     TestResults(failed=1, attempted=...)
+
+    >>> traceback._COLORIZE = save_colorize
     """
     exc = ValueError('Text')
     exc.add_note(note)
@@ -3419,6 +3459,9 @@ def test_syntax_error_subclass_from_stdlib():
 
 def test_syntax_error_with_incorrect_expected_note():
     """
+    >>> save_colorize = traceback._COLORIZE
+    >>> traceback._COLORIZE = False
+
     >>> def f(x):
     ...     r'''
     ...     >>> exc = SyntaxError("error", ("x.py", 23, None, "bad syntax"))
@@ -3447,6 +3490,8 @@ def test_syntax_error_with_incorrect_expected_note():
         note1
         note2
     TestResults(failed=1, attempted=...)
+
+    >>> traceback._COLORIZE = save_colorize
     """
 
 
diff --git a/Lib/traceback.py b/Lib/traceback.py
@@ -448,8 +448,12 @@ class _ANSIColors:
     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."""
diff --git a/Misc/NEWS.d/next/Library/2024-04-06-18-41-36.gh-issue-117225.tJh1Hw.rst b/Misc/NEWS.d/next/Library/2024-04-06-18-41-36.gh-issue-117225.tJh1Hw.rst
@@ -0,0 +1 @@
+Add colour to doctest output. Patch by Hugo van Kemenade.

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+Add colour to doctest output. Patch by Hugo van Kemenade.`