Add _interpreters.capture_exception().

python · ericsnowcurrently · Apr 11, 2024 · Apr 8, 2024 · Apr 5, 2024 · Apr 5, 2024
commit c75a115c3e6c2f71594fa5191465a075518fb55a
@@ -217,6 +217,11 @@ typedef struct _excinfo {
     const char *errdisplay;
 } _PyXI_excinfo;
 
+PyAPI_FUNC(int) _PyXI_InitExcInfo(_PyXI_excinfo *info, PyObject *exc);
+PyAPI_FUNC(PyObject *) _PyXI_FormatExcInfo(_PyXI_excinfo *info);
+PyAPI_FUNC(PyObject *) _PyXI_ExcInfoAsObject(_PyXI_excinfo *info);
+PyAPI_FUNC(void) _PyXI_ClearExcInfo(_PyXI_excinfo *info);
+
 
 typedef enum error_code {
     _PyXI_ERR_NO_ERROR = 0,

diff --git a/Modules/_xxsubinterpretersmodule.c b/Modules/_xxsubinterpretersmodule.c
@@ -1199,9 +1199,78 @@ interp_decref(PyObject *self, PyObject *args, PyObject *kwds)
 }
 
 
+static PyObject *
+capture_exception(PyObject *self, PyObject *args, PyObject *kwds)
+{
+    static char *kwlist[] = {"exc", NULL};
+    PyObject *exc_arg = NULL;
+    if (!PyArg_ParseTupleAndKeywords(args, kwds,
+                                     "|O:capture_exception", kwlist,
+                                     &exc_arg))
+    {
+        return NULL;
+    }
+
+    PyObject *exc = exc_arg;
+    if (exc == NULL || exc == Py_None) {
+        exc = PyErr_GetRaisedException();
+        if (exc == NULL) {
+            Py_RETURN_NONE;
+        }
+    }
+    else if (!PyExceptionInstance_Check(exc)) {
+        PyErr_Format(PyExc_TypeError, "expected exception, got %R", exc);
+        return NULL;
+    }
+    PyObject *captured = NULL;
+
+    _PyXI_excinfo info = {0};
+    if (_PyXI_InitExcInfo(&info, exc) < 0) {
+        goto finally;
+    }
+    captured = _PyXI_ExcInfoAsObject(&info);
+    if (captured == NULL) {
+        goto finally;
+    }
+
+    PyObject *formatted = _PyXI_FormatExcInfo(&info);
+    if (formatted == NULL) {
+        Py_CLEAR(captured);
+        goto finally;
+    }
+    int res = PyObject_SetAttrString(captured, "formatted", formatted);
+    Py_DECREF(formatted);
+    if (res < 0) {
+        Py_CLEAR(captured);
+        goto finally;
+    }
+
+finally:
+    _PyXI_ClearExcInfo(&info);
+    if (exc != exc_arg) {
+        if (PyErr_Occurred()) {
+            PyErr_SetRaisedException(exc);
+        }
+        else {
+            _PyErr_ChainExceptions1(exc);
+        }
+    }
+    return captured;
+}
+
+PyDoc_STRVAR(capture_exception_doc,
+"capture_exception(exc=None) -> types.SimpleNamespace\n\
+\n\
+Return a snapshot of an exception.  If \"exc\" is None\n\
+then the current exception, if any, is used (but not cleared).\n\
+\n\
+The returned snapshot is the same as what _interpreters.exec() returns.");
+
+
 static PyMethodDef module_functions[] = {
     {"new_config",                _PyCFunction_CAST(interp_new_config),
      METH_VARARGS | METH_KEYWORDS, new_config_doc},
+
     {"create",                    _PyCFunction_CAST(interp_create),
      METH_VARARGS | METH_KEYWORDS, create_doc},
     {"destroy",                   _PyCFunction_CAST(interp_destroy),
@@ -1228,14 +1297,18 @@ static PyMethodDef module_functions[] = {
 
     {"set___main___attrs",        _PyCFunction_CAST(interp_set___main___attrs),
      METH_VARARGS, set___main___attrs_doc},
-    {"is_shareable",              _PyCFunction_CAST(object_is_shareable),
-     METH_VARARGS | METH_KEYWORDS, is_shareable_doc},
 
     {"incref",                    _PyCFunction_CAST(interp_incref),
      METH_VARARGS | METH_KEYWORDS, NULL},
     {"decref",                    _PyCFunction_CAST(interp_decref),
      METH_VARARGS | METH_KEYWORDS, NULL},
 
+    {"is_shareable",              _PyCFunction_CAST(object_is_shareable),
+     METH_VARARGS | METH_KEYWORDS, is_shareable_doc},
+
+    {"capture_exception",         _PyCFunction_CAST(capture_exception),
+     METH_VARARGS | METH_KEYWORDS, capture_exception_doc},
+
     {NULL,                        NULL}           /* sentinel */
 };
 

@@ -468,7 +468,7 @@ _release_xid_data(_PyCrossInterpreterData *data, int rawfree)
 /***********************/
 
 static int
-_excinfo_init_type(struct _excinfo_type *info, PyObject *exc)
+_excinfo_init_type_from_exception(struct _excinfo_type *info, PyObject *exc)
 {
     /* Note that this copies directly rather than into an intermediate
        struct and does not clear on error.  If we need that then we
@@ -504,7 +504,7 @@ _excinfo_init_type(struct _excinfo_type *info, PyObject *exc)
     }
     info->qualname = _copy_string_obj_raw(strobj, NULL);
     Py_DECREF(strobj);
-    if (info->name == NULL) {
+    if (info->qualname == NULL) {
         return -1;
     }
 
@@ -515,10 +515,51 @@ _excinfo_init_type(struct _excinfo_type *info, PyObject *exc)
     }
     info->module = _copy_string_obj_raw(strobj, NULL);
     Py_DECREF(strobj);
+    if (info->module == NULL) {
+        return -1;
+    }
+
+    return 0;
+}
+
+static int
+_excinfo_init_type_from_object(struct _excinfo_type *info, PyObject *exctype)
+{
+    PyObject *strobj = NULL;
+
+    // __name__
+    strobj = PyObject_GetAttrString(exctype, "__name__");
+    if (strobj == NULL) {
+        return -1;
+    }
+    info->name = _copy_string_obj_raw(strobj, NULL);
+    Py_DECREF(strobj);
     if (info->name == NULL) {
         return -1;
     }
 
+    // __qualname__
+    strobj = PyObject_GetAttrString(exctype, "__qualname__");
+    if (strobj == NULL) {
+        return -1;
+    }
+    info->qualname = _copy_string_obj_raw(strobj, NULL);
+    Py_DECREF(strobj);
+    if (info->qualname == NULL) {
+        return -1;
+    }
+
+    // __module__
+    strobj = PyObject_GetAttrString(exctype, "__module__");
+    if (strobj == NULL) {
+        return -1;
+    }
+    info->module = _copy_string_obj_raw(strobj, NULL);
+    Py_DECREF(strobj);
+    if (info->module == NULL) {
+        return -1;
+    }
+
     return 0;
 }
 
@@ -584,7 +625,7 @@ _PyXI_excinfo_Clear(_PyXI_excinfo *info)
     *info = (_PyXI_excinfo){{NULL}};
 }
 
-static PyObject *
+PyObject *
 _PyXI_excinfo_format(_PyXI_excinfo *info)
 {
     const char *module, *qualname;
@@ -627,7 +668,7 @@ _PyXI_excinfo_InitFromException(_PyXI_excinfo *info, PyObject *exc)
     }
     const char *failure = NULL;
 
-    if (_excinfo_init_type(&info->type, exc) < 0) {
+    if (_excinfo_init_type_from_exception(&info->type, exc) < 0) {
         failure = "error while initializing exception type snapshot";
         goto error;
     }
@@ -672,6 +713,57 @@ _PyXI_excinfo_InitFromException(_PyXI_excinfo *info, PyObject *exc)
     return failure;
 }
 
+static const char *
+_PyXI_excinfo_InitFromObject(_PyXI_excinfo *info, PyObject *obj)
+{
+    const char *failure = NULL;
+
+    PyObject *exctype = PyObject_GetAttrString(obj, "type");
+    if (exctype == NULL) {
+        failure = "exception snapshot missing 'type' attribute";
+        goto error;
+    }
+    int res = _excinfo_init_type_from_object(&info->type, exctype);
+    Py_DECREF(exctype);
+    if (res < 0) {
+        failure = "error while initializing exception type snapshot";
+        goto error;
+    }
+
+    // Extract the exception message.
+    PyObject *msgobj = PyObject_GetAttrString(obj, "msg");
+    if (msgobj == NULL) {
+        failure = "exception snapshot missing 'msg' attribute";
+        goto error;
+    }
+    info->msg = _copy_string_obj_raw(msgobj, NULL);
+    Py_DECREF(msgobj);
+    if (info->msg == NULL) {
+        failure = "error while copying exception message";
+        goto error;
+    }
+
+    // Pickle a traceback.TracebackException.
+    PyObject *errdisplay = PyObject_GetAttrString(obj, "errdisplay");
+    if (errdisplay == NULL) {
+        failure = "exception snapshot missing 'errdisplay' attribute";
+        goto error;
+    }
+    info->errdisplay = _copy_string_obj_raw(errdisplay, NULL);
+    Py_DECREF(errdisplay);
+    if (info->errdisplay == NULL) {
+        failure = "error while copying exception error display";
+        goto error;
+    }
+
+    return NULL;
+
+error:
+    assert(failure != NULL);
+    _PyXI_excinfo_Clear(info);
+    return failure;
+}
+
 static void
 _PyXI_excinfo_Apply(_PyXI_excinfo *info, PyObject *exctype)
 {
@@ -825,6 +917,47 @@ _PyXI_excinfo_AsObject(_PyXI_excinfo *info)
 }
 
 
+int
+_PyXI_InitExcInfo(_PyXI_excinfo *info, PyObject *exc)
+{
+    assert(!PyErr_Occurred());
+    if (exc == NULL || exc == Py_None) {
+        PyErr_SetString(PyExc_ValueError, "missing exc");
+        return -1;
+    }
+    const char *failure;
+    if (PyExceptionInstance_Check(exc) || PyExceptionClass_Check(exc)) {
+        failure = _PyXI_excinfo_InitFromException(info, exc);
+    }
+    else {
+        failure = _PyXI_excinfo_InitFromObject(info, exc);
+    }
+    if (failure != NULL) {
+        PyErr_SetString(PyExc_Exception, failure);
+        return -1;
+    }
+    return 0;
+}
+
+PyObject *
+_PyXI_FormatExcInfo(_PyXI_excinfo *info)
+{
+    return _PyXI_excinfo_format(info);
+}
+
+PyObject *
+_PyXI_ExcInfoAsObject(_PyXI_excinfo *info)
+{
+    return _PyXI_excinfo_AsObject(info);
+}
+
+void
+_PyXI_ClearExcInfo(_PyXI_excinfo *info)
+{
+    _PyXI_excinfo_Clear(info);
+}
+
+
 /***************************/
 /* short-term data sharing */
 /***************************/