bpo-45607: Make it possible to enrich exception displays via setting their __note__ field (GH-29880)

1 files changed, 35 insertions, 0 deletions

diff --git a/Python/pythonrun.c b/Python/pythonrun.c
index 2f68b21460..5a118b4821 100644
--- a/Python/pythonrun.c
+++ b/Python/pythonrun.c
@@ -1083,6 +1083,41 @@ print_exception(struct exception_print_context *ctx, PyObject *value)
         PyErr_Clear();
     }
     err += PyFile_WriteString("\n", f);
+
+    if (err == 0 && PyExceptionInstance_Check(value)) {
+        _Py_IDENTIFIER(__note__);
+
+        PyObject *note = _PyObject_GetAttrId(value, &PyId___note__);
+        if (note == NULL) {
+            err = -1;
+        }
+        if (err == 0 && PyUnicode_Check(note)) {
+            _Py_static_string(PyId_newline, "\n");
+            PyObject *lines = PyUnicode_Split(
+                note, _PyUnicode_FromId(&PyId_newline), -1);
+            if (lines == NULL) {
+                err = -1;
+            }
+            else {
+                Py_ssize_t n = PyList_GET_SIZE(lines);
+                for (Py_ssize_t i = 0; i < n; i++) {
+                    if (err == 0) {
+                        PyObject *line = PyList_GET_ITEM(lines, i);
+                        assert(PyUnicode_Check(line));
+                        err = write_indented_margin(ctx, f);
+                        if (err == 0) {
+                            err = PyFile_WriteObject(line, f, Py_PRINT_RAW);
+                        }
+                        if (err == 0) {
+                            err = PyFile_WriteString("\n", f);
+                        }
+                    }
+                }
+            }
+            Py_DECREF(lines);
+        }
+        Py_XDECREF(note);
+    }
     Py_XDECREF(tb);
     Py_DECREF(value);
     /* If an error happened here, don't show it.