bpo-43950: support some multi-line expressions for PEP 657

Lib/test/test_doctest.py

@@ -2835,6 +2835,7 @@ def test_unicode(): """
         Traceback (most recent call last):
           File ...
             exec(compile(example.source, filename, "single",
+            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
           File "<doctest foo-bär@baz[0]>", line 1, in <module>
             raise Exception('clé')
             ^^^^^^^^^^^^^^^^^^^^^^

Lib/test/test_traceback.py

@@ -429,6 +429,30 @@ def f_with_multiline():
             '    ^^^^^^^^^^\n'
             f'  File "{__file__}", line {lineno_f+1}, in f_with_multiline\n'
             '    raise ValueError(\n'
+            '    ^^^^^^^^^^^^^^^^^'
+        )
+        result_lines = self.get_exception(f_with_multiline)
+        self.assertEqual(result_lines, expected_f.splitlines())
+
+    def test_caret_multiline_expression_bin_op(self):
+        # Make sure no carets are printed for expressions spanning multiple
+        # lines.
+        def f_with_multiline():
+            return (
+                1 /
+                0 +
+                2
+            )
+
+        lineno_f = f_with_multiline.__code__.co_firstlineno
+        expected_f = (
+            '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_multiline\n'
+            '    1 /\n'
+            '    ^^^'
         )
         result_lines = self.get_exception(f_with_multiline)
         self.assertEqual(result_lines, expected_f.splitlines())

Lib/traceback.py

@@ -4,6 +4,7 @@
 import itertools
 import linecache
 import sys
+from contextlib import suppress
 
 __all__ = ['extract_stack', 'extract_tb', 'format_exception',
            'format_exception_only', 'format_list', 'format_stack',
@@ -463,19 +464,20 @@ def format_frame(self, frame):
 
             stripped_characters = len(frame._original_line) - len(frame.line.lstrip())
             if (
-                frame.end_lineno == frame.lineno
-                and frame.colno is not None
+                frame.colno is not None
                 and frame.end_colno is not None
             ):
                 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 - 1]
-                    )
-                except Exception:
-                    anchors = None
+                anchors = None
+                if frame.lineno == frame.end_lineno:
+                    with suppress(Exception):
+                        anchors = _extract_caret_anchors_from_line_segment(
+                            frame._original_line[colno - 1:end_colno - 1]
+                        )
+                else:
+                    end_colno = stripped_characters + len(frame.line.strip())
 
                 row.append('    ')
                 row.append(' ' * (colno - stripped_characters))

Python/traceback.c

@@ -720,11 +720,11 @@ tb_displayline(PyTracebackObject* tb, PyObject *f, PyObject *filename, int linen
                               &end_line, &end_col_byte_offset)) {
         goto done;
     }
-    if (start_line != end_line) {
-        goto done;
-    }
 
-    if (start_col_byte_offset < 0 || end_col_byte_offset < 0) {
+    if (start_line < 0 || end_line < 0
+        || start_col_byte_offset < 0
+        || end_col_byte_offset < 0)
+    {
         goto done;
     }
 
@@ -762,11 +762,30 @@ tb_displayline(PyTracebackObject* tb, PyObject *f, PyObject *filename, int linen
     char *primary_error_char = "^";
     char *secondary_error_char = primary_error_char;
 
-    int res = extract_anchors_from_line(filename, source_line, start_offset, end_offset,
-                                        &left_end_offset, &right_start_offset,
-                                        &primary_error_char, &secondary_error_char);
-    if (res < 0 && ignore_source_errors() < 0) {
-        goto done;
+    if (start_line == end_line) {
+        int res = extract_anchors_from_line(filename, source_line, start_offset, end_offset,
+                                            &left_end_offset, &right_start_offset,
+                                            &primary_error_char, &secondary_error_char);
+        if (res < 0 && ignore_source_errors() < 0) {
+            goto done;
+        }
+    }
+    else {
+        // If this is a multi-line expression, then we will highlight until
+        // the last non-whitespace character.
+        const char *source_line_str = PyUnicode_AsUTF8(source_line);
+        if (!source_line_str) {
+            goto done;
+        }
+
+        Py_ssize_t i = PyUnicode_GET_LENGTH(source_line);
+        while (--i >= 0) {
+            if (!IS_WHITESPACE(source_line_str[i])) {
+                break;
+            }
+        }
+
+        end_offset = i + 1;
     }
 
     err = print_error_location_carets(f, truncation, start_offset, end_offset,