bpo-45292: [PEP 654] Update traceback display code to work with exception groups (GH-29207)

1 files changed, 93 insertions, 25 deletions

diff --git a/Python/traceback.c b/Python/traceback.c
index 22a0922..67f995a 100644
--- a/Python/traceback.c
+++ b/Python/traceback.c
@@ -14,6 +14,7 @@
 #include "pycore_pyarena.h"       // _PyArena_Free()
 #include "pycore_pyerrors.h"      // _PyErr_Fetch()
 #include "pycore_pystate.h"       // _PyThreadState_GET()
+#include "pycore_traceback.h"     // EXCEPTION_TB_HEADER
 #include "../Parser/pegen.h"      // _PyPegen_byte_offset_to_character_offset()
 #include "structmember.h"         // PyMemberDef
 #include "osdefs.h"               // SEP
@@ -379,8 +380,44 @@ finally:
     return result;
 }
 
+/* Writes indent spaces. Returns 0 on success and non-zero on failure.
+ */
+int
+_Py_WriteIndent(int indent, PyObject *f)
+{
+    int err = 0;
+    char buf[11] = "          ";
+    assert(strlen(buf) == 10);
+    while (indent > 0) {
+        if (indent < 10) {
+            buf[indent] = '\0';
+        }
+        err = PyFile_WriteString(buf, f);
+        if (err != 0) {
+            return err;
+        }
+        indent -= 10;
+    }
+    return 0;
+}
+
+/* Writes indent spaces, followed by the margin if it is not `\0`.
+   Returns 0 on success and non-zero on failure.
+ */
 int
-_Py_DisplaySourceLine(PyObject *f, PyObject *filename, int lineno, int indent, int *truncation, PyObject **line)
+_Py_WriteIndentedMargin(int indent, const char *margin, PyObject *f)
+{
+    int err = _Py_WriteIndent(indent, f);
+    if (err == 0 && margin) {
+        err = PyFile_WriteString(margin, f);
+    }
+    return err;
+}
+
+static int
+display_source_line_with_margin(PyObject *f, PyObject *filename, int lineno, int indent,
+                                int margin_indent, const char *margin,
+                                int *truncation, PyObject **line)
 {
     int err = 0;
     int fd;
@@ -508,27 +545,33 @@ _Py_DisplaySourceLine(PyObject *f, PyObject *filename, int lineno, int indent, i
         *truncation = i - indent;
     }
 
+    if (err == 0) {
+        err = _Py_WriteIndentedMargin(margin_indent, margin, f);
+    }
     /* Write some spaces before the line */
-    strcpy(buf, "          ");
-    assert (strlen(buf) == 10);
-    while (indent > 0) {
-        if (indent < 10)
-            buf[indent] = '\0';
-        err = PyFile_WriteString(buf, f);
-        if (err != 0)
-            break;
-        indent -= 10;
+    if (err == 0) {
+        err = _Py_WriteIndent(indent, f);
     }
 
     /* finally display the line */
-    if (err == 0)
+    if (err == 0) {
         err = PyFile_WriteObject(lineobj, f, Py_PRINT_RAW);
+    }
     Py_DECREF(lineobj);
-    if  (err == 0)
+    if  (err == 0) {
         err = PyFile_WriteString("\n", f);
+    }
     return err;
 }
 
+int
+_Py_DisplaySourceLine(PyObject *f, PyObject *filename, int lineno, int indent,
+                      int *truncation, PyObject **line)
+{
+    return display_source_line_with_margin(f, filename, lineno, indent, 0,
+                                           NULL, truncation, line);
+}
+
 /* AST based Traceback Specialization
  *
  * When displaying a new traceback line, for certain syntactical constructs
@@ -697,7 +740,7 @@ print_error_location_carets(PyObject *f, int offset, Py_ssize_t start_offset, Py
 
 static int
 tb_displayline(PyTracebackObject* tb, PyObject *f, PyObject *filename, int lineno,
-               PyFrameObject *frame, PyObject *name)
+               PyFrameObject *frame, PyObject *name, int margin_indent, const char *margin)
 {
     int err;
     PyObject *line;
@@ -708,15 +751,20 @@ tb_displayline(PyTracebackObject* tb, PyObject *f, PyObject *filename, int linen
                                 filename, lineno, name);
     if (line == NULL)
         return -1;
-    err = PyFile_WriteObject(line, f, Py_PRINT_RAW);
+    err = _Py_WriteIndentedMargin(margin_indent, margin, f);
+    if (err == 0) {
+        err = PyFile_WriteObject(line, f, Py_PRINT_RAW);
+    }
     Py_DECREF(line);
     if (err != 0)
         return err;
 
     int truncation = _TRACEBACK_SOURCE_LINE_INDENT;
     PyObject* source_line = NULL;
-    if (_Py_DisplaySourceLine(f, filename, lineno, _TRACEBACK_SOURCE_LINE_INDENT,
-                               &truncation, &source_line) != 0 || !source_line) {
+    int rc = display_source_line_with_margin(
+            f, filename, lineno, _TRACEBACK_SOURCE_LINE_INDENT,
+            margin_indent, margin, &truncation, &source_line);
+    if (rc != 0 || !source_line) {
         /* ignore errors since we can't report them, can we? */
         err = ignore_source_errors();
         goto done;
@@ -801,9 +849,12 @@ tb_displayline(PyTracebackObject* tb, PyObject *f, PyObject *filename, int linen
         end_offset = i + 1;
     }
 
-    err = print_error_location_carets(f, truncation, start_offset, end_offset,
-                                      right_start_offset, left_end_offset,
-                                      primary_error_char, secondary_error_char);
+    err = _Py_WriteIndentedMargin(margin_indent, margin, f);
+    if (err == 0) {
+        err = print_error_location_carets(f, truncation, start_offset, end_offset,
+                                          right_start_offset, left_end_offset,
+                                          primary_error_char, secondary_error_char);
+    }
 
 done:
     Py_XDECREF(source_line);
@@ -830,7 +881,8 @@ tb_print_line_repeated(PyObject *f, long cnt)
 }
 
 static int
-tb_printinternal(PyTracebackObject *tb, PyObject *f, long limit)
+tb_printinternal(PyTracebackObject *tb, PyObject *f, long limit,
+                 int indent, const char *margin)
 {
     int err = 0;
     Py_ssize_t depth = 0;
@@ -864,7 +916,7 @@ tb_printinternal(PyTracebackObject *tb, PyObject *f, long limit)
         cnt++;
         if (err == 0 && cnt <= TB_RECURSIVE_CUTOFF) {
             err = tb_displayline(tb, f, code->co_filename, tb->tb_lineno,
-                                 tb->tb_frame, code->co_name);
+                                 tb->tb_frame, code->co_name, indent, margin);
             if (err == 0) {
                 err = PyErr_CheckSignals();
             }
@@ -881,7 +933,8 @@ tb_printinternal(PyTracebackObject *tb, PyObject *f, long limit)
 #define PyTraceBack_LIMIT 1000
 
 int
-PyTraceBack_Print(PyObject *v, PyObject *f)
+_PyTraceBack_Print_Indented(PyObject *v, int indent, const char *margin,
+                            const char *header_margin, const char *header, PyObject *f)
 {
     int err;
     PyObject *limitv;
@@ -904,12 +957,27 @@ PyTraceBack_Print(PyObject *v, PyObject *f)
             return 0;
         }
     }
-    err = PyFile_WriteString("Traceback (most recent call last):\n", f);
-    if (!err)
-        err = tb_printinternal((PyTracebackObject *)v, f, limit);
+    err = _Py_WriteIndentedMargin(indent, header_margin, f);
+    if (err == 0) {
+        err = PyFile_WriteString(header, f);
+    }
+    if (err == 0) {
+        err = tb_printinternal((PyTracebackObject *)v, f, limit, indent, margin);
+    }
     return err;
 }
 
+int
+PyTraceBack_Print(PyObject *v, PyObject *f)
+{
+    int indent = 0;
+    const char *margin = NULL;
+    const char *header_margin = NULL;
+    const char *header = EXCEPTION_TB_HEADER;
+
+    return _PyTraceBack_Print_Indented(v, indent, margin, header_margin, header, f);
+}
+
 /* Format an integer in range [0; 0xffffffff] to decimal and write it
    into the file fd.