Port new specializations to traceback.py

ammaraskar · isidentical · isidentical · commit d79d365a8b08 · 2021-07-10T12:06:39.000+03:00
Co-authored-by: Batuhan Taskaya &lt;batuhanosmantaskaya@gmail.com&gt;
diff --git a/Doc/library/traceback.rst b/Doc/library/traceback.rst
@@ -473,7 +473,7 @@ The output for the example would look similar to this:
    ['Traceback (most recent call last):\n',
     '  File "<doctest default[0]>", line 10, in <module>\n    lumberjack()\n    ^^^^^^^^^^^^\n',
     '  File "<doctest default[0]>", line 4, in lumberjack\n    bright_side_of_death()\n    ^^^^^^^^^^^^^^^^^^^^^^\n',
-    '  File "<doctest default[0]>", line 7, in bright_side_of_death\n    return tuple()[0]\n           ^^^^^^^^^^\n',
+    '  File "<doctest default[0]>", line 7, in bright_side_of_death\n    return tuple()[0]\n           ~~~~~~~^^^\n',
     'IndexError: tuple index out of range\n']
    *** extract_tb:
    [<FrameSummary file <doctest...>, line 10 in <module>>,
@@ -482,7 +482,7 @@ The output for the example would look similar to this:
    *** format_tb:
    ['  File "<doctest default[0]>", line 10, in <module>\n    lumberjack()\n    ^^^^^^^^^^^^\n',
     '  File "<doctest default[0]>", line 4, in lumberjack\n    bright_side_of_death()\n    ^^^^^^^^^^^^^^^^^^^^^^\n',
-    '  File "<doctest default[0]>", line 7, in bright_side_of_death\n    return tuple()[0]\n           ^^^^^^^^^^\n']
+    '  File "<doctest default[0]>", line 7, in bright_side_of_death\n    return tuple()[0]\n           ~~~~~~~^^^\n']
    *** tb_lineno: 10
 
 
diff --git a/Lib/test/test_traceback.py b/Lib/test/test_traceback.py
@@ -406,6 +406,61 @@ def f_with_multiline():
         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())
+
+
 
 @cpython_only
 @requires_debug_ranges()
@@ -1615,7 +1670,7 @@ def f():
         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
@@ -494,9 +494,23 @@ def format(self):
                     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[0]))
+                        row.append('^' * (anchors[1] - anchors[0]))
+                        row.append('~' * (end_colno - colno - anchors[1]))
+                    else:
+                        row.append('^' * (end_colno - colno))
+
                     row.append('\n')
 
             if frame.locals:
@@ -520,6 +534,39 @@ def _byte_offset_to_character_offset(str, offset):
     return len(as_utf8[:offset + 1].decode("utf-8"))
 
 
+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 left_anchor, right_anchor
+                case ast.Subscript():
+                    return expr.value.end_col_offset, expr.slice.end_col_offset + 1
+
+    return None
+
+
 class TracebackException:
     """An exception ready for rendering.
 
