From 7b2f4b198db68507c1da3e3d7f1628f7b11f98e9 Mon Sep 17 00:00:00 2001
From: Zak Stucke <zakstucke@hotmail.co.uk>
Date: Wed, 19 Jul 2023 11:58:24 +0300
Subject: [PATCH] Prevent traceback loss on conversion to and from PyErr

---
 newsfragments/3328.fixed.md |  2 ++
 src/err/mod.rs              | 19 ++++++++++++---
 src/types/traceback.rs      | 47 ++++++++++++++++++++++++++++++++++++-
 3 files changed, 64 insertions(+), 4 deletions(-)
 create mode 100644 newsfragments/3328.fixed.md

diff --git a/newsfragments/3328.fixed.md b/newsfragments/3328.fixed.md
new file mode 100644
index 00000000000..70f4c1cf8f4
--- /dev/null
+++ b/newsfragments/3328.fixed.md
@@ -0,0 +1,2 @@
+Fixed `PyErr.from_value()` to maintain original traceback.
+Fixed `PyErr.from_value()` to maintain original traceback.
diff --git a/src/err/mod.rs b/src/err/mod.rs
index 8516cc9b1db..1cfd403564e 100644
--- a/src/err/mod.rs
+++ b/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.
diff --git a/src/types/traceback.rs b/src/types/traceback.rs
index 36a04e599d1..6b6dd3122d4 100644
--- a/src/types/traceback.rs
+++ b/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));
+        })
+    }
 }