gh-112730: Use color to highlight error locations (gh-112732)

Signed-off-by: Pablo Galindo <pablogsal@gmail.com>
Co-authored-by: Łukasz Langa <lukasz@langa.pl>

Author: pablogsal

Doc/using/cmdline.rst

@@ -612,6 +612,27 @@ Miscellaneous options
    .. versionadded:: 3.13
       The ``-X presite`` option.
 
+Controlling Color
+~~~~~~~~~~~~~~~~~
+
+The Python interpreter is configured by default to use colors to highlight
+output in certain situations such as when displaying tracebacks. This
+behavior can be controlled by setting different environment variables.
+
+Setting the environment variable ``TERM`` to ``dumb`` will disable color.
+
+If the environment variable ``FORCE_COLOR`` is set, then color will be
+enabled regardless of the value of TERM. This is useful on CI systems which
+aren’t terminals but can none-the-less display ANSI escape sequences.
+
+If the environment variable ``NO_COLOR`` is set, Python will disable all color
+in the output. This takes precedence over ``FORCE_COLOR``.
+
+All these environment variables are used also by other tools to control color
+output. To control the color output only in the Python interpreter, the
+:envvar:`PYTHON_COLORS` environment variable can be used. This variable takes
+precedence over ``NO_COLOR``, which in turn takes precedence over
+``FORCE_COLOR``.
 
 Options you shouldn't use
 ~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1110,6 +1131,12 @@ conflict.
 
    .. versionadded:: 3.13
 
+.. envvar:: PYTHON_COLORS
+
+   If this variable is set to ``1``, the interpreter will colorize various kinds
+   of output. Setting it to ``0`` deactivates this behavior.
+
+   .. versionadded:: 3.13
 
 Debug-mode variables
 ~~~~~~~~~~~~~~~~~~~~

Doc/whatsnew/3.13.rst

@@ -85,7 +85,13 @@ Important deprecations, removals or restrictions:
 New Features
 ============
 
+Improved Error Messages
+-----------------------
 
+* The interpreter now colorizes error messages when displaying tracebacks by default.
+  This feature can be controlled via the new :envvar:`PYTHON_COLORS` environment
+  variable as well as the canonical ``NO_COLOR`` and ``FORCE_COLOR`` environment
+  variables. (Contributed by Pablo Galindo Salgado in :gh:`112730`.)
 
 Other Language Changes
 ======================

Lib/test/test_traceback.py

@@ -8,6 +8,7 @@
 import inspect
 import builtins
 import unittest
+import unittest.mock
 import re
 import tempfile
 import random
@@ -24,6 +25,7 @@
 import json
 import textwrap
 import traceback
+import contextlib
 from functools import partial
 from pathlib import Path
 
@@ -41,6 +43,14 @@
 class TracebackCases(unittest.TestCase):
     # For now, a very minimal set of tests.  I want to be sure that
     # formatting of SyntaxErrors works based on changes for 2.1.
+    def setUp(self):
+        super().setUp()
+        self.colorize = traceback._COLORIZE
+        traceback._COLORIZE = False
+
+    def tearDown(self):
+        super().tearDown()
+        traceback._COLORIZE = self.colorize
 
     def get_exception_format(self, func, exc):
         try:
@@ -521,7 +531,7 @@ def test_signatures(self):
         self.assertEqual(
             str(inspect.signature(traceback.print_exception)),
             ('(exc, /, value=<implicit>, tb=<implicit>, '
-             'limit=None, file=None, chain=True)'))
+             'limit=None, file=None, chain=True, **kwargs)'))
 
         self.assertEqual(
             str(inspect.signature(traceback.format_exception)),
@@ -3031,7 +3041,7 @@ def some_inner(k, v):
 
     def test_custom_format_frame(self):
         class CustomStackSummary(traceback.StackSummary):
-            def format_frame_summary(self, frame_summary):
+            def format_frame_summary(self, frame_summary, colorize=False):
                 return f'{frame_summary.filename}:{frame_summary.lineno}'
 
         def some_inner():
@@ -3056,7 +3066,7 @@ def g():
         tb = g()
 
         class Skip_G(traceback.StackSummary):
-            def format_frame_summary(self, frame_summary):
+            def format_frame_summary(self, frame_summary, colorize=False):
                 if frame_summary.name == 'g':
                     return None
                 return super().format_frame_summary(frame_summary)
@@ -3076,7 +3086,6 @@ def __repr__(self) -> str:
         raise Exception("Unrepresentable")
 
 class TestTracebackException(unittest.TestCase):
-
     def do_test_smoke(self, exc, expected_type_str):
         try:
             raise exc
@@ -4245,6 +4254,115 @@ def test_levenshtein_distance_short_circuit(self):
                 res3 = traceback._levenshtein_distance(a, b, threshold)
                 self.assertGreater(res3, threshold, msg=(a, b, threshold))
 
+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 = traceback._ANSIColors.RED
+        boldr = traceback._ANSIColors.BOLD_RED
+        reset = traceback._ANSIColors.RESET
+        self.assertIn("y = " + red + "x['a']['b']" + reset + boldr + "['c']" + reset, lines)
+        self.assertIn("return " + red + "foo" + reset + boldr + "(1,2,3,4)" + reset, lines)
+        self.assertIn("return " + red + "baz" + reset + boldr + "(1," + reset, lines)
+        self.assertIn(boldr + "2,3" + reset, lines)
+        self.assertIn(boldr + ",4)" + reset, lines)
+        self.assertIn(red + "bar" + reset + boldr + "()" + reset, 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))
+        red = traceback._ANSIColors.RED
+        magenta = traceback._ANSIColors.MAGENTA
+        boldm = traceback._ANSIColors.BOLD_MAGENTA
+        boldr = traceback._ANSIColors.BOLD_RED
+        reset = traceback._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('traceback._can_colorize', return_value=True):
+                    exception_print(e)
+            actual = tbstderr.getvalue().splitlines()
+
+        red = traceback._ANSIColors.RED
+        boldr = traceback._ANSIColors.BOLD_RED
+        magenta = traceback._ANSIColors.MAGENTA
+        boldm = traceback._ANSIColors.BOLD_MAGENTA
+        reset = traceback._ANSIColors.RESET
+        lno_foo = foo.__code__.co_firstlineno
+        expected = ['Traceback (most recent call last):',
+            f'  File {magenta}"{__file__}"{reset}, '
+            f'line {magenta}{lno_foo+5}{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()
+        with virtual_patching:
+            with unittest.mock.patch("os.isatty") as isatty_mock:
+                isatty_mock.return_value = True
+                with unittest.mock.patch("os.environ", {'TERM': 'dumb'}):
+                    self.assertEqual(traceback._can_colorize(), False)
+                with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '1'}):
+                    self.assertEqual(traceback._can_colorize(), True)
+                with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '0'}):
+                    self.assertEqual(traceback._can_colorize(), False)
+                with unittest.mock.patch("os.environ", {'NO_COLOR': '1'}):
+                    self.assertEqual(traceback._can_colorize(), False)
+                with unittest.mock.patch("os.environ", {'NO_COLOR': '1', "PYTHON_COLORS": '1'}):
+                    self.assertEqual(traceback._can_colorize(), True)
+                with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1'}):
+                    self.assertEqual(traceback._can_colorize(), True)
+                with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1', 'NO_COLOR': '1'}):
+                    self.assertEqual(traceback._can_colorize(), False)
+                with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1', "PYTHON_COLORS": '0'}):
+                    self.assertEqual(traceback._can_colorize(), False)
+                isatty_mock.return_value = False
+                self.assertEqual(traceback._can_colorize(), False)
 
 if __name__ == "__main__":
     unittest.main()