use PyErr_SetRaisedException in raise_lazy

src/err/err_state.rs

@@ -6,6 +6,8 @@ use std::{
 
 #[cfg(not(Py_3_12))]
 use crate::sync::MutexExt;
+#[cfg(Py_3_12)]
+use crate::types::{PyString, PyTuple};
 use crate::{
     exceptions::{PyBaseException, PyTypeError},
     ffi,
@@ -383,29 +385,141 @@ fn lazy_into_normalized_ffi_tuple(
 }
 
 /// Raises a "lazy" exception state into the Python interpreter.
-///
-/// In principle this could be split in two; first a function to create an exception
-/// in a normalized state, and then a call to `PyErr_SetRaisedException` to raise it.
-///
-/// This would require either moving some logic from C to Rust, or requesting a new
-/// API in CPython.
 fn raise_lazy(py: Python<'_>, lazy: Box<PyErrStateLazyFn>) {
     let PyErrStateLazyFnOutput { ptype, pvalue } = lazy(py);
+
     unsafe {
+        #[cfg(not(Py_3_12))]
         if ffi::PyExceptionClass_Check(ptype.as_ptr()) == 0 {
             ffi::PyErr_SetString(
                 PyTypeError::type_object_raw(py).cast(),
                 c"exceptions must derive from BaseException".as_ptr(),
-            )
+            );
         } else {
-            ffi::PyErr_SetObject(ptype.as_ptr(), pvalue.as_ptr())
+            ffi::PyErr_SetObject(ptype.as_ptr(), pvalue.as_ptr());
+        }
+
+        #[cfg(Py_3_12)]
+        {
+            let exc = create_normalized_exception(ptype.bind(py), pvalue.into_bound(py));
+
+            ffi::PyErr_SetRaisedException(exc.into_ptr());
+        }
+    }
+}
+
+#[cfg(Py_3_12)]
+fn create_normalized_exception<'py>(
+    ptype: &Bound<'py, PyAny>,
+    mut pvalue: Bound<'py, PyAny>,
+) -> Bound<'py, PyBaseException> {
+    let py = ptype.py();
+
+    // 1: check type is a subclass of BaseException
+    let ptype: Bound<'py, PyType> = if unsafe { ffi::PyExceptionClass_Check(ptype.as_ptr()) } == 0 {
+        pvalue = PyString::new(py, "exceptions must derive from BaseException").into_any();
+        PyTypeError::type_object(py)
+    } else {
+        // Safety: PyExceptionClass_Check guarantees that ptype is a subclass of BaseException
+        unsafe { ptype.cast_unchecked() }.clone()
+    };
+
+    let current_handled_exception: Option<Bound<'_, PyBaseException>> = unsafe {
+        ffi::PyErr_GetHandledException()
+            .assume_owned_or_opt(py)
+            .map(|obj| obj.cast_into_unchecked())
+    };
+
+    let pvalue = if pvalue.is_exact_instance(&ptype) {
+        // Safety: already an exception value of the correct type
+        let exc = unsafe { pvalue.cast_into_unchecked::<PyBaseException>() };
+
+        if let Some(current_context) = current_handled_exception.as_ref() {
+            // Check if this exception is already in the context chain, so we do not create reference cycles in the context chain.
+            for obj in context_chain_iter(current_context.clone()) {
+                if obj.is(&exc) {
+                    // Loop in context chain, breaking the loop
+                    unsafe { ffi::PyException_SetContext(obj.as_ptr(), std::ptr::null_mut()) };
+                    break;
+                }
+            }
+        }
+
+        Ok(exc)
+    } else if pvalue.is_none() {
+        // None -> no arguments
+        ptype.call0().and_then(|pvalue| Ok(pvalue.cast_into()?))
+    } else if let Ok(tup) = pvalue.cast::<PyTuple>() {
+        // Tuple -> use as tuple of arguments
+        ptype.call1(tup).and_then(|pvalue| Ok(pvalue.cast_into()?))
+    } else {
+        // Anything else -> use as single argument
+        ptype
+            .call1((pvalue,))
+            .and_then(|pvalue| Ok(pvalue.cast_into()?))
+    };
+
+    match pvalue {
+        Ok(pvalue) => {
+            // Implicitly set the context of the new exception to the currently handled exception, if any.
+            if let Some(context) = current_handled_exception {
+                unsafe { ffi::PyException_SetContext(pvalue.as_ptr(), context.into_ptr()) };
+            }
+            pvalue
+        }
+        Err(e) => e.value(py).clone(),
+    }
+}
+
+/// Iterates through the context chain of exceptions, starting from `start`, and yields each exception in the chain.
+/// The iterator is resilient to loops in the context chain, and will terminate if a loop is detected.
+#[cfg(Py_3_12)]
+fn context_chain_iter(
+    start: Bound<'_, PyBaseException>,
+) -> impl Iterator<Item = Bound<'_, PyBaseException>> {
+    let mut slow = None;
+    let mut current = Some(start);
+    let mut slow_update_toggle = false;
+
+    #[inline]
+    fn get_next<'py>(current: &Bound<'py, PyBaseException>) -> Option<Bound<'py, PyBaseException>> {
+        unsafe {
+            ffi::PyException_GetContext(current.as_ptr())
+                .assume_owned_or_opt(current.py())
+                .map(|obj| obj.cast_into_unchecked())
         }
     }
+
+    std::iter::from_fn(move || {
+        let next = get_next(current.as_ref()?);
+
+        // Detect loops in the context chain using Floyd's Tortoise and Hare algorithm.
+        if let Some((current_slow, current_fast)) = slow.as_ref().zip(next.as_ref()) {
+            if current_fast.is(current_slow) {
+                // Loop detected
+                current.take();
+                return None;
+            }
+
+            // Every second iteration, advance the slow pointer by one step
+            if slow_update_toggle {
+                slow = get_next(current_slow);
+            }
+
+            slow_update_toggle = !slow_update_toggle;
+        }
+
+        // Set the slow pointer after the first iteration
+        if slow.is_none() {
+            slow = current.clone()
+        }
+
+        std::mem::replace(&mut current, next)
+    })
 }
 
 #[cfg(test)]
 mod tests {
-
     use crate::{
         exceptions::PyValueError, sync::PyOnceLock, Py, PyAny, PyErr, PyErrArguments, Python,
     };
@@ -478,4 +592,113 @@ mod tests {
                 .is_instance_of::<PyValueError>(py))
         });
     }
+
+    #[test]
+    #[cfg(feature = "macros")]
+    fn test_new_exception_context() {
+        use crate::{
+            exceptions::{PyRuntimeError, PyValueError},
+            pyfunction,
+            types::{PyDict, PyDictMethods},
+            wrap_pyfunction, PyResult,
+        };
+        #[pyfunction(crate = "crate")]
+        fn throw_exception() -> PyResult<()> {
+            Err(PyValueError::new_err("error happened"))
+        }
+
+        Python::attach(|py| {
+            let globals = PyDict::new(py);
+            let f = wrap_pyfunction!(throw_exception, py).unwrap();
+            globals.set_item("throw_exception", f).unwrap();
+            let err = py
+                .run(
+                    c"try:\n  raise RuntimeError()\nexcept:\n  throw_exception()\n",
+                    Some(&globals),
+                    None,
+                )
+                .unwrap_err();
+
+            let context = err.context(py).unwrap();
+            assert!(context.is_instance_of::<PyRuntimeError>(py))
+        })
+    }
+
+    #[test]
+    #[cfg(Py_3_12)]
+    fn compare_create_normalized_exception_with_pyerr_setobject() {
+        use crate::{
+            conversion::IntoPyObjectExt, err::err_state::PyErrStateNormalized,
+            exceptions::PyRuntimeError, ffi, type_object::PyTypeInfo, types::any::PyAnyMethods,
+            Bound,
+        };
+
+        fn test_exception<'py>(ptype: &Bound<'py, PyAny>, pvalue: Bound<'py, PyAny>) {
+            let py = ptype.py();
+
+            let exc1 = super::create_normalized_exception(ptype, pvalue.clone());
+
+            unsafe {
+                ffi::PyErr_SetObject(ptype.as_ptr(), pvalue.as_ptr());
+            }
+            let exc2 = PyErrStateNormalized::take(py)
+                .unwrap()
+                .pvalue
+                .into_bound(py);
+
+            let err1 = PyErr::from_value(exc1.into_any());
+            let err2 = PyErr::from_value(exc2.into_any());
+
+            assert!(err1.get_type(py).is(err2.get_type(py)));
+            assert!(err1.context(py).xor(err2.context(py)).is_none());
+            assert!(err1.traceback(py).xor(err2.traceback(py)).is_none());
+            assert!(err1.cause(py).xor(err2.cause(py)).is_none());
+            assert_eq!(err1.to_string(), err2.to_string());
+        }
+
+        Python::attach(|py| {
+            test_exception(&PyRuntimeError::type_object(py), py.None().into_bound(py));
+
+            test_exception(
+                &PyRuntimeError::type_object(py),
+                "Boom".into_bound_py_any(py).unwrap(),
+            );
+
+            test_exception(
+                &PyRuntimeError::type_object(py),
+                (3, 2, 1, "Boom").into_bound_py_any(py).unwrap(),
+            );
+
+            test_exception(
+                &PyRuntimeError::type_object(py),
+                PyRuntimeError::new_err("Boom")
+                    .into_value(py)
+                    .into_any()
+                    .into_bound(py),
+            );
+        })
+    }
+
+    #[test]
+    #[cfg(Py_3_12)]
+    fn test_context_chain_iter_terminates() {
+        fn create_loop(py: Python<'_>, size: usize) -> PyErr {
+            let first = PyValueError::new_err("exc0");
+            let last = (1..size).fold(first.clone_ref(py), |prev, i| {
+                let exc = PyValueError::new_err(format!("exc{i}"));
+                prev.set_context(py, Some(exc.clone_ref(py)));
+                exc
+            });
+            last.set_context(py, Some(first.clone_ref(py)));
+
+            last
+        }
+
+        Python::attach(|py| {
+            for size in 1..=8 {
+                let chain = create_loop(py, size);
+                super::context_chain_iter(chain.into_value(py).into_bound(py)).for_each(drop);
+            }
+        })
+    }
 }