Add PyErr::set_context, PyErr::context

Author: bschoenmaeckers

newsfragments/5887.added.md

@@ -0,0 +1 @@
+Add `PyErr::set_context`, `PyErr::context`, `ffi::PyErr_GetHandledException` and `ffi::PyErr_SetHandledException`.

pyo3-ffi/src/pyerrors.rs

@@ -53,6 +53,12 @@ extern_libpython! {
     pub fn PyErr_GetRaisedException() -> *mut PyObject;
     #[cfg(Py_3_12)]
     pub fn PyErr_SetRaisedException(exc: *mut PyObject);
+    #[cfg(Py_3_11)]
+    #[cfg_attr(PyPy, link_name = "PyPyErr_GetHandledException")]
+    pub fn PyErr_GetHandledException() -> *mut PyObject;
+    #[cfg(Py_3_11)]
+    #[cfg_attr(PyPy, link_name = "PyPyErr_SetHandledException")]
+    pub fn PyErr_SetHandledException(exc: *mut PyObject);
     #[cfg_attr(PyPy, link_name = "PyPyException_SetTraceback")]
     pub fn PyException_SetTraceback(arg1: *mut PyObject, arg2: *mut PyObject) -> c_int;
     #[cfg_attr(PyPy, link_name = "PyPyException_GetTraceback")]

src/err/mod.rs

@@ -604,6 +604,29 @@ impl PyErr {
         }
     }
 
+    /// Return the context (either an exception instance, or None, set by an implicit exception
+    /// during handling of another exception) associated with the exception, as accessible from
+    /// Python through `__context__`.
+    pub fn context(&self, py: Python<'_>) -> Option<PyErr> {
+        unsafe {
+            ffi::PyException_GetContext(self.value(py).as_ptr())
+                .assume_owned_or_opt(py)
+                .map(Self::from_value)
+        }
+    }
+
+    /// Set the context associated with the exception, pass `None` to clear it.
+    pub fn set_context(&self, py: Python<'_>, context: Option<PyErr>) {
+        let value = self.value(py);
+        let context = context.map(|err| err.into_value(py));
+        unsafe {
+            ffi::PyException_SetContext(
+                value.as_ptr(),
+                context.map_or(std::ptr::null_mut(), Py::into_ptr),
+            );
+        }
+    }
+
     /// Equivalent to calling `add_note` on the exception in Python.
     #[cfg(Py_3_11)]
     pub fn add_note<N: for<'py> IntoPyObject<'py, Target = PyString>>(
@@ -1027,4 +1050,19 @@ mod tests {
             );
         });
     }
+
+    #[test]
+    fn test_set_context() {
+        Python::attach(|py| {
+            let err = PyErr::new::<PyValueError, _>("original error");
+            assert!(err.context(py).is_none());
+
+            let context = PyErr::new::<PyTypeError, _>("context error");
+            err.set_context(py, Some(context));
+            assert!(err.context(py).unwrap().is_instance_of::<PyTypeError>(py));
+
+            err.set_context(py, None);
+            assert!(err.context(py).is_none());
+        })
+    }
 }