bpo-43950: Specialize tracebacks for subscripts/binary ops (GH-27037)

Co-authored-by: Ammar Askar <ammar@ammaraskar.com>
Co-authored-by: Pablo Galindo <pablogsal@gmail.com>

2 files changed, 139 insertions, 3 deletions

diff --git a/Lib/test/test_traceback.py b/Lib/test/test_traceback.py
index 50ebccef82a..8baf38c1afd 100644
--- a/Lib/test/test_traceback.py
+++ b/Lib/test/test_traceback.py
@@ -12,9 +12,11 @@ from test.support import (Error, captured_output, cpython_only, ALWAYS_EQ,
                           requires_debug_ranges, has_no_debug_ranges)
 from test.support.os_helper import TESTFN, unlink
 from test.support.script_helper import assert_python_ok, assert_python_failure
-import textwrap
 
+import os
+import textwrap
 import traceback
+from functools import partial
 
 
 test_code = namedtuple('code', ['co_filename', 'co_name'])
@@ -406,6 +408,82 @@ class TracebackErrorLocationCaretTests(unittest.TestCase):
         result_lines = self.get_exception(f_with_multiline)
         self.assertEqual(result_lines, expected_f.splitlines())
 
+    def test_caret_for_binary_operators(self):
+        def f_with_binary_operator():
+            divisor = 20
+            return 10 + divisor / 0 + 30
+
+        lineno_f = f_with_binary_operator.__code__.co_firstlineno
+        expected_error = (
+            'Traceback (most recent call last):\n'
+            f'  File "{__file__}", line {self.callable_line}, in get_exception\n'
+            '    callable()\n'
+            '    ^^^^^^^^^^\n'
+            f'  File "{__file__}", line {lineno_f+2}, in f_with_binary_operator\n'
+            '    return 10 + divisor / 0 + 30\n'
+            '                ~~~~~~~~^~~\n'
+        )
+        result_lines = self.get_exception(f_with_binary_operator)
+        self.assertEqual(result_lines, expected_error.splitlines())
+
+    def test_caret_for_binary_operators_two_char(self):
+        def f_with_binary_operator():
+            divisor = 20
+            return 10 + divisor // 0 + 30
+
+        lineno_f = f_with_binary_operator.__code__.co_firstlineno
+        expected_error = (
+            'Traceback (most recent call last):\n'
+            f'  File "{__file__}", line {self.callable_line}, in get_exception\n'
+            '    callable()\n'
+            '    ^^^^^^^^^^\n'
+            f'  File "{__file__}", line {lineno_f+2}, in f_with_binary_operator\n'
+            '    return 10 + divisor // 0 + 30\n'
+            '                ~~~~~~~~^^~~\n'
+        )
+        result_lines = self.get_exception(f_with_binary_operator)
+        self.assertEqual(result_lines, expected_error.splitlines())
+
+    def test_caret_for_subscript(self):
+        def f_with_subscript():
+            some_dict = {'x': {'y': None}}
+            return some_dict['x']['y']['z']
+
+        lineno_f = f_with_subscript.__code__.co_firstlineno
+        expected_error = (
+            'Traceback (most recent call last):\n'
+            f'  File "{__file__}", line {self.callable_line}, in get_exception\n'
+            '    callable()\n'
+            '    ^^^^^^^^^^\n'
+            f'  File "{__file__}", line {lineno_f+2}, in f_with_subscript\n'
+            "    return some_dict['x']['y']['z']\n"
+            '           ~~~~~~~~~~~~~~~~~~~^^^^^\n'
+        )
+        result_lines = self.get_exception(f_with_subscript)
+        self.assertEqual(result_lines, expected_error.splitlines())
+
+    def test_traceback_specialization_with_syntax_error(self):
+        bytecode = compile("1 / 0 / 1 / 2\n", TESTFN, "exec")
+
+        with open(TESTFN, "w") as file:
+            # make the file's contents invalid
+            file.write("1 $ 0 / 1 / 2\n")
+        self.addCleanup(unlink, TESTFN)
+
+        func = partial(exec, bytecode)
+        result_lines = self.get_exception(func)
+
+        lineno_f = bytecode.co_firstlineno
+        expected_error = (
+            'Traceback (most recent call last):\n'
+            f'  File "{__file__}", line {self.callable_line}, in get_exception\n'
+            '    callable()\n'
+            '    ^^^^^^^^^^\n'
+            f'  File "{TESTFN}", line {lineno_f}, in <module>\n'
+            "    1 $ 0 / 1 / 2\n"
+            '    ^^^^^\n'
+        )
+        self.assertEqual(result_lines, expected_error.splitlines())
 
 @cpython_only
 @requires_debug_ranges()
@@ -1615,7 +1693,7 @@ class TestTracebackException(unittest.TestCase):
         self.assertEqual(
             output.getvalue().split('\n')[-5:],
             ['    x/0',
-             '    ^^^',
+             '    ~^~',
              '    x = 12',
              'ZeroDivisionError: division by zero',
              ''])
diff --git a/Lib/traceback.py b/Lib/traceback.py
index 7cb124188ac..ec5e20d431f 100644
--- a/Lib/traceback.py
+++ b/Lib/traceback.py
@@ -494,9 +494,23 @@ class StackSummary(list):
                     colno = _byte_offset_to_character_offset(frame._original_line, frame.colno)
                     end_colno = _byte_offset_to_character_offset(frame._original_line, frame.end_colno)
 
+                    try:
+                        anchors = _extract_caret_anchors_from_line_segment(
+                            frame._original_line[colno - 1:end_colno]
+                        )
+                    except Exception:
+                        anchors = None
+
                     row.append('    ')
                     row.append(' ' * (colno - stripped_characters))
-                    row.append('^' * (end_colno - colno))
+
+                    if anchors:
+                        row.append(anchors.primary_char * (anchors.left_end_offset))
+                        row.append(anchors.secondary_char * (anchors.right_start_offset - anchors.left_end_offset))
+                        row.append(anchors.primary_char * (end_colno - colno - anchors.right_start_offset))
+                    else:
+                        row.append('^' * (end_colno - colno))
+
                     row.append('\n')
 
             if frame.locals:
@@ -520,6 +534,50 @@ def _byte_offset_to_character_offset(str, offset):
     return len(as_utf8[:offset + 1].decode("utf-8"))
 
 
+_Anchors = collections.namedtuple(
+    "_Anchors",
+    [
+        "left_end_offset",
+        "right_start_offset",
+        "primary_char",
+        "secondary_char",
+    ],
+    defaults=["~", "^"]
+)
+
+def _extract_caret_anchors_from_line_segment(segment):
+    import ast
+
+    try:
+        tree = ast.parse(segment)
+    except SyntaxError:
+        return None
+
+    if len(tree.body) != 1:
+        return None
+
+    statement = tree.body[0]
+    match statement:
+        case ast.Expr(expr):
+            match expr:
+                case ast.BinOp():
+                    operator_str = segment[expr.left.end_col_offset:expr.right.col_offset]
+                    operator_offset = len(operator_str) - len(operator_str.lstrip())
+
+                    left_anchor = expr.left.end_col_offset + operator_offset
+                    right_anchor = left_anchor + 1
+                    if (
+                        operator_offset + 1 < len(operator_str)
+                        and not operator_str[operator_offset + 1].isspace()
+                    ):
+                        right_anchor += 1
+                    return _Anchors(left_anchor, right_anchor)
+                case ast.Subscript():
+                    return _Anchors(expr.value.end_col_offset, expr.slice.end_col_offset + 1)
+
+    return None
+
+
 class TracebackException:
     """An exception ready for rendering.