Merge pull request #3328 from zakstucke/main

Prevent traceback loss on conversion to and from PyErr

newsfragments/3328.fixed.md

@@ -0,0 +1 @@
+Fixed `PyErr.from_value()` and `PyErr.into_value()` to both maintain original traceback instead of losing it on conversion.

src/err/mod.rs

@@ -177,10 +177,16 @@ impl PyErr {
     /// ```
     pub fn from_value(obj: &PyAny) -> PyErr {
         let state = if let Ok(obj) = obj.downcast::<PyBaseException>() {
+            let pvalue: Py<PyBaseException> = obj.into();
+
+            let ptraceback = unsafe {
+                Py::from_owned_ptr_or_opt(obj.py(), ffi::PyException_GetTraceback(pvalue.as_ptr()))
+            };
+
             PyErrState::Normalized(PyErrStateNormalized {
                 ptype: obj.get_type().into(),
-                pvalue: obj.into(),
-                ptraceback: None,
+                pvalue,
+                ptraceback,
             })
         } else {
             // Assume obj is Type[Exception]; let later normalization handle if this
@@ -228,7 +234,14 @@ impl PyErr {
         // NB technically this causes one reference count increase and decrease in quick succession
         // on pvalue, but it's probably not worth optimizing this right now for the additional code
         // complexity.
-        self.normalized(py).pvalue.clone_ref(py)
+        let normalized = self.normalized(py);
+        let exc = normalized.pvalue.clone_ref(py);
+        if let Some(tb) = normalized.ptraceback.as_ref() {
+            unsafe {
+                ffi::PyException_SetTraceback(exc.as_ptr(), tb.as_ptr());
+            }
+        }
+        exc
     }
 
     /// Returns the traceback of this exception object.

src/types/traceback.rs

@@ -65,7 +65,7 @@ impl PyTraceback {
 
 #[cfg(test)]
 mod tests {
-    use crate::Python;
+    use crate::{prelude::*, types::PyDict};
 
     #[test]
     fn format_traceback() {
@@ -80,4 +80,49 @@ mod tests {
             );
         })
     }
+
+    #[test]
+    fn test_err_from_value() {
+        Python::with_gil(|py| {
+            let locals = PyDict::new(py);
+            // Produce an error from python so that it has a traceback
+            py.run(
+                r"
+try:
+    raise ValueError('raised exception')
+except Exception as e:
+    err = e
+",
+                None,
+                Some(locals),
+            )
+            .unwrap();
+            let err = PyErr::from_value(locals.get_item("err").unwrap());
+            let traceback = err.value(py).getattr("__traceback__").unwrap();
+            assert!(err.traceback(py).unwrap().is(traceback));
+        })
+    }
+
+    #[test]
+    fn test_err_into_py() {
+        Python::with_gil(|py| {
+            let locals = PyDict::new(py);
+            // Produce an error from python so that it has a traceback
+            py.run(
+                r"
+def f():
+    raise ValueError('raised exception')
+",
+                None,
+                Some(locals),
+            )
+            .unwrap();
+            let f = locals.get_item("f").unwrap();
+            let err = f.call0().unwrap_err();
+            let traceback = err.traceback(py).unwrap();
+            let err_object = err.clone_ref(py).into_py(py).into_ref(py);
+
+            assert!(err_object.getattr("__traceback__").unwrap().is(traceback));
+        })
+    }
 }