eagerly normalize fetched exception on Python 3.11 and older

Author: davidhewitt

newsfragments/4655.changed.md

@@ -0,0 +1 @@
+Eagerly normalize exceptions in `PyErr::take()` and `PyErr::fetch()` on Python 3.11 and older.

src/err/err_state.rs

@@ -3,7 +3,8 @@ use std::cell::UnsafeCell;
 use crate::{
     exceptions::{PyBaseException, PyTypeError},
     ffi,
-    types::{PyTraceback, PyType},
+    ffi_ptr_ext::FfiPtrExt,
+    types::{PyAnyMethods, PyTraceback, PyType},
     Bound, Py, PyAny, PyErrArguments, PyObject, PyTypeInfo, Python,
 };
 
@@ -34,19 +35,6 @@ impl PyErrState {
         })))
     }
 
-    #[cfg(not(Py_3_12))]
-    pub(crate) fn ffi_tuple(
-        ptype: PyObject,
-        pvalue: Option<PyObject>,
-        ptraceback: Option<PyObject>,
-    ) -> Self {
-        Self::from_inner(PyErrStateInner::FfiTuple {
-            ptype,
-            pvalue,
-            ptraceback,
-        })
-    }
-
     pub(crate) fn normalized(normalized: PyErrStateNormalized) -> Self {
         Self::from_inner(PyErrStateInner::Normalized(normalized))
     }
@@ -151,19 +139,61 @@ impl PyErrStateNormalized {
 
     #[cfg(Py_3_12)]
     pub(crate) fn ptraceback<'py>(&self, py: Python<'py>) -> Option<Bound<'py, PyTraceback>> {
-        use crate::ffi_ptr_ext::FfiPtrExt;
-        use crate::types::any::PyAnyMethods;
         unsafe {
             ffi::PyException_GetTraceback(self.pvalue.as_ptr())
                 .assume_owned_or_opt(py)
                 .map(|b| b.downcast_into_unchecked())
         }
     }
 
-    #[cfg(Py_3_12)]
     pub(crate) fn take(py: Python<'_>) -> Option<PyErrStateNormalized> {
-        unsafe { Py::from_owned_ptr_or_opt(py, ffi::PyErr_GetRaisedException()) }
-            .map(|pvalue| PyErrStateNormalized { pvalue })
+        #[cfg(Py_3_12)]
+        {
+            // Safety: PyErr_GetRaisedException can be called when attached to Python and
+            // returns either NULL or an owned reference.
+            unsafe { ffi::PyErr_GetRaisedException().assume_owned_or_opt(py) }.map(|pvalue| {
+                PyErrStateNormalized {
+                    // Safety: PyErr_GetRaisedException returns a valid exception type.
+                    pvalue: unsafe { pvalue.downcast_into_unchecked() }.unbind(),
+                }
+            })
+        }
+
+        #[cfg(not(Py_3_12))]
+        {
+            let (ptype, pvalue, ptraceback) = unsafe {
+                let mut ptype: *mut ffi::PyObject = std::ptr::null_mut();
+                let mut pvalue: *mut ffi::PyObject = std::ptr::null_mut();
+                let mut ptraceback: *mut ffi::PyObject = std::ptr::null_mut();
+
+                ffi::PyErr_Fetch(&mut ptype, &mut pvalue, &mut ptraceback);
+
+                // Ensure that the exception coming from the interpreter is normalized.
+                if !ptype.is_null() {
+                    ffi::PyErr_NormalizeException(&mut ptype, &mut pvalue, &mut ptraceback);
+                }
+
+                // Safety: PyErr_NormalizeException will have produced up to three owned
+                // references of the correct types.
+                (
+                    ptype
+                        .assume_owned_or_opt(py)
+                        .map(|b| b.downcast_into_unchecked()),
+                    pvalue
+                        .assume_owned_or_opt(py)
+                        .map(|b| b.downcast_into_unchecked()),
+                    ptraceback
+                        .assume_owned_or_opt(py)
+                        .map(|b| b.downcast_into_unchecked()),
+                )
+            };
+
+            ptype.map(|ptype| PyErrStateNormalized {
+                ptype: ptype.unbind(),
+                pvalue: pvalue.expect("normalized exception value missing").unbind(),
+                ptraceback: ptraceback.map(Bound::unbind),
+            })
+        }
     }
 
     #[cfg(not(Py_3_12))]
@@ -204,12 +234,6 @@ pub(crate) type PyErrStateLazyFn =
 
 enum PyErrStateInner {
     Lazy(Box<PyErrStateLazyFn>),
-    #[cfg(not(Py_3_12))]
-    FfiTuple {
-        ptype: PyObject,
-        pvalue: Option<PyObject>,
-        ptraceback: Option<PyObject>,
-    },
     Normalized(PyErrStateNormalized),
 }
 
@@ -231,20 +255,6 @@ impl PyErrStateInner {
                 PyErrStateNormalized::take(py)
                     .expect("exception missing after writing to the interpreter")
             }
-            #[cfg(not(Py_3_12))]
-            PyErrStateInner::FfiTuple {
-                ptype,
-                pvalue,
-                ptraceback,
-            } => {
-                let mut ptype = ptype.into_ptr();
-                let mut pvalue = pvalue.map_or(std::ptr::null_mut(), Py::into_ptr);
-                let mut ptraceback = ptraceback.map_or(std::ptr::null_mut(), Py::into_ptr);
-                unsafe {
-                    ffi::PyErr_NormalizeException(&mut ptype, &mut pvalue, &mut ptraceback);
-                    PyErrStateNormalized::from_normalized_ffi_tuple(py, ptype, pvalue, ptraceback)
-                }
-            }
             PyErrStateInner::Normalized(normalized) => normalized,
         }
     }
@@ -253,15 +263,6 @@ impl PyErrStateInner {
     fn restore(self, py: Python<'_>) {
         let (ptype, pvalue, ptraceback) = match self {
             PyErrStateInner::Lazy(lazy) => lazy_into_normalized_ffi_tuple(py, lazy),
-            PyErrStateInner::FfiTuple {
-                ptype,
-                pvalue,
-                ptraceback,
-            } => (
-                ptype.into_ptr(),
-                pvalue.map_or(std::ptr::null_mut(), Py::into_ptr),
-                ptraceback.map_or(std::ptr::null_mut(), Py::into_ptr),
-            ),
             PyErrStateInner::Normalized(PyErrStateNormalized {
                 ptype,
                 pvalue,

src/err/mod.rs

@@ -355,60 +355,6 @@ impl PyErr {
     /// expected to have been set, for example from [`PyErr::occurred`] or by an error return value
     /// from a C FFI function, use [`PyErr::fetch`].
     pub fn take(py: Python<'_>) -> Option<PyErr> {
-        Self::_take(py)
-    }
-
-    #[cfg(not(Py_3_12))]
-    fn _take(py: Python<'_>) -> Option<PyErr> {
-        let (ptype, pvalue, ptraceback) = unsafe {
-            let mut ptype: *mut ffi::PyObject = std::ptr::null_mut();
-            let mut pvalue: *mut ffi::PyObject = std::ptr::null_mut();
-            let mut ptraceback: *mut ffi::PyObject = std::ptr::null_mut();
-            ffi::PyErr_Fetch(&mut ptype, &mut pvalue, &mut ptraceback);
-
-            // Convert to Py immediately so that any references are freed by early return.
-            let ptype = PyObject::from_owned_ptr_or_opt(py, ptype);
-            let pvalue = PyObject::from_owned_ptr_or_opt(py, pvalue);
-            let ptraceback = PyObject::from_owned_ptr_or_opt(py, ptraceback);
-
-            // A valid exception state should always have a non-null ptype, but the other two may be
-            // null.
-            let ptype = match ptype {
-                Some(ptype) => ptype,
-                None => {
-                    debug_assert!(
-                        pvalue.is_none(),
-                        "Exception type was null but value was not null"
-                    );
-                    debug_assert!(
-                        ptraceback.is_none(),
-                        "Exception type was null but traceback was not null"
-                    );
-                    return None;
-                }
-            };
-
-            (ptype, pvalue, ptraceback)
-        };
-
-        if ptype.as_ptr() == PanicException::type_object_raw(py).cast() {
-            let msg = pvalue
-                .as_ref()
-                .and_then(|obj| obj.bind(py).str().ok())
-                .map(|py_str| py_str.to_string_lossy().into())
-                .unwrap_or_else(|| String::from("Unwrapped panic from Python code"));
-
-            let state = PyErrState::ffi_tuple(ptype, pvalue, ptraceback);
-            Self::print_panic_and_unwind(py, state, msg)
-        }
-
-        Some(PyErr::from_state(PyErrState::ffi_tuple(
-            ptype, pvalue, ptraceback,
-        )))
-    }
-
-    #[cfg(Py_3_12)]
-    fn _take(py: Python<'_>) -> Option<PyErr> {
         let state = PyErrStateNormalized::take(py)?;
         let pvalue = state.pvalue.bind(py);
         if pvalue.get_type().as_ptr() == PanicException::type_object_raw(py).cast() {