Inject rust stack frames into traceback

Author: bschoenmaeckers

Cargo.toml

@@ -61,6 +61,7 @@ uuid = { version = "1.12.0", optional = true }
 lock_api = { version = "0.4", optional = true }
 parking_lot = { version = "0.12", optional = true }
 iana-time-zone = { version = "0.1", optional = true, features = ["fallback"]}
+backtrace = "0.3.76"
 
 [target.'cfg(not(target_has_atomic = "64"))'.dependencies]
 portable-atomic = "1.0"
@@ -80,6 +81,7 @@ tempfile = "3.12.0"
 static_assertions = "1.1.0"
 uuid = { version = "1.10.0", features = ["v4"] }
 parking_lot = { version = "0.12.3", features = ["arc_lock"] }
+insta = { version = "1.46.3", features = ["filters"] }
 
 [build-dependencies]
 pyo3-build-config = { path = "pyo3-build-config", version = "=0.28.2", features = ["resolve-config"] }

src/err/err_state.rs

@@ -4,6 +4,8 @@ use std::{
     thread::ThreadId,
 };
 
+#[cfg(debug_assertions)]
+use crate::err::backtrace_to_frames;
 #[cfg(not(Py_3_12))]
 use crate::sync::MutexExt;
 use crate::{
@@ -36,10 +38,15 @@ impl PyErrState {
     }
 
     pub(crate) fn lazy_arguments(ptype: Py<PyAny>, args: impl PyErrArguments + 'static) -> Self {
+        #[cfg(debug_assertions)]
+        let backtrace = backtrace::Backtrace::new_unresolved();
+
         Self::from_inner(PyErrStateInner::Lazy(Box::new(move |py| {
             PyErrStateLazyFnOutput {
                 ptype,
                 pvalue: args.arguments(py),
+                #[cfg(debug_assertions)]
+                backtrace,
             }
         })))
     }
@@ -301,6 +308,8 @@ impl PyErrStateNormalized {
 pub(crate) struct PyErrStateLazyFnOutput {
     pub(crate) ptype: Py<PyAny>,
     pub(crate) pvalue: Py<PyAny>,
+    #[cfg(debug_assertions)]
+    pub(crate) backtrace: backtrace::Backtrace,
 }
 
 pub(crate) type PyErrStateLazyFn =
@@ -390,15 +399,28 @@ fn lazy_into_normalized_ffi_tuple(
 /// 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);
+    let PyErrStateLazyFnOutput {
+        ptype,
+        pvalue,
+        #[cfg(debug_assertions)]
+        mut backtrace,
+    } = lazy(py);
+
     unsafe {
         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(debug_assertions)]
+            {
+                for frame in backtrace_to_frames(py, &mut backtrace) {
+                    ffi::PyTraceBack_Here(frame.into_ptr().cast());
+                }
+            }
         }
     }
 }

src/err/mod.rs

@@ -12,14 +12,14 @@ use crate::types::any::PyAnyMethods;
 #[cfg(Py_3_11)]
 use crate::types::PyString;
 use crate::types::{
-    string::PyStringMethods, traceback::PyTracebackMethods, typeobject::PyTypeMethods, PyTraceback,
-    PyType,
+    string::PyStringMethods, traceback::PyTracebackMethods, typeobject::PyTypeMethods, PyFrame,
+    PyTraceback, PyType,
 };
 use crate::{exceptions::PyBaseException, ffi};
 use crate::{BoundObject, Py, PyAny, Python};
 use err_state::{PyErrState, PyErrStateLazyFnOutput, PyErrStateNormalized};
 use std::convert::Infallible;
-use std::ffi::CStr;
+use std::ffi::{CStr, CString};
 
 mod cast_error;
 mod downcast_error;
@@ -127,10 +127,14 @@ impl PyErr {
         T: PyTypeInfo,
         A: PyErrArguments + Send + Sync + 'static,
     {
+        #[cfg(debug_assertions)]
+        let backtrace = backtrace::Backtrace::new_unresolved();
         PyErr::from_state(PyErrState::lazy(Box::new(move |py| {
             PyErrStateLazyFnOutput {
                 ptype: T::type_object(py).into(),
                 pvalue: args.arguments(py),
+                #[cfg(debug_assertions)]
+                backtrace,
             }
         })))
     }
@@ -289,7 +293,24 @@ impl PyErr {
             Self::print_panic_and_unwind(py, state)
         }
 
-        Some(PyErr::from_state(PyErrState::normalized(state)))
+        let err = PyErr::from_state(PyErrState::normalized(state));
+
+        #[cfg(debug_assertions)]
+        {
+            let mut backtrace = backtrace::Backtrace::new();
+            if let Some(traceback) =
+                PyTraceback::from_frames(py, backtrace_to_frames(py, &mut backtrace))
+                    .ok()
+                    .flatten()
+            {
+                if let Some(prev) = err.traceback(py) {
+                    let _ = traceback.append(prev);
+                }
+                err.set_traceback(py, Some(traceback));
+            }
+        }
+
+        Some(err)
     }
 
     #[cold]
@@ -696,6 +717,53 @@ impl<'py> IntoPyObject<'py> for PyErr {
     }
 }
 
+#[cfg(debug_assertions)]
+fn backtrace_to_frames<'py, 'a>(
+    py: Python<'py>,
+    backtrace: &'a mut backtrace::Backtrace,
+) -> impl Iterator<Item = Bound<'py, PyFrame>> + use<'py, 'a> {
+    backtrace.resolve();
+    backtrace
+        .frames()
+        .iter()
+        .flat_map(|frame| frame.symbols())
+        .take_while(|&symbol| {
+            symbol
+                .name()
+                .map(|sym| sym.to_string())
+                .map(|ref name| {
+                    dbg!(name);
+                    !(name.starts_with("pyo3::impl_::trampoline::")
+                        || name.contains("__rust_begin_short_backtrace"))
+                })
+                .unwrap_or(true)
+        })
+        .filter_map(move |symbol| {
+            let file = CString::new(
+                symbol
+                    .filename()
+                    .unwrap_or_else(|| std::path::Path::new("<unknown>"))
+                    .as_os_str()
+                    .to_string_lossy()
+                    .as_ref(),
+            )
+            .ok()?;
+
+            let function = CString::new(
+                symbol
+                    .name()
+                    .unwrap()
+                    .to_string()
+                    .trim_end_matches(char::is_alphanumeric)
+                    .trim_end_matches(':'),
+            )
+            .ok()?;
+            let line = symbol.lineno().unwrap_or(0);
+
+            PyFrame::new(py, &file, &function, line as _).ok()
+        })
+}
+
 impl<'py> IntoPyObject<'py> for &PyErr {
     type Target = PyBaseException;
     type Output = Bound<'py, Self::Target>;

tests/snapshots/test_backtrace__new_err_function_err_result@macos.snap

@@ -0,0 +1,58 @@
+---
+source: tests/test_backtrace.rs
+expression: traceback
+---
+Traceback (most recent call last):
+  File "[RUST_CORE]/ops/function.rs", line 250, in core::ops::function::FnOnce::call_once
+  File "[RUST_CORE]/ops/function.rs", line 250, in core::ops::function::FnOnce::call_once
+    extern "rust-call" fn call_once(self, args: Args) -> Self::Output;
+  File "./tests/test_backtrace.rs", line 7, in test_backtrace::test_new_err_function_err_result::{{closure}}
+    fn test_new_err_function_err_result() {
+  File "./tests/test_backtrace.rs", line 16, in test_backtrace::test_new_err_function_err_result
+    Python::attach(|py| {
+  File "./src/marker.rs", line 415, in pyo3::marker::Python::attach
+    f(guard.python())
+  File "./tests/test_backtrace.rs", line 22, in test_backtrace::test_new_err_function_err_result::{{closure}}
+    .run(
+  File "./src/marker.rs", line 641, in pyo3::marker::Python::run
+    code.run(globals, locals).map(|obj| {
+  File "./src/types/code.rs", line 139, in <pyo3::instance::Bound<pyo3::types::code::PyCode> as pyo3::types::code::PyCodeMethods>::run
+    .assume_owned_or_err(self.py())
+  File "./src/ffi_ptr_ext.rs", line 43, in <*mut pyo3_ffi::object::PyObject as pyo3::ffi_ptr_ext::FfiPtrExt>::assume_owned_or_err
+    unsafe { Bound::from_owned_ptr_or_err(py, self) }
+  File "./src/instance.rs", line 385, in pyo3::instance::Bound<pyo3::types::any::PyAny>::from_owned_ptr_or_err
+    None => Err(PyErr::fetch(py)),
+  File "./src/err/mod.rs", line 352, in pyo3::err::PyErr::fetch
+    PyErr::take(py).unwrap_or_else(failed_to_fetch)
+  File "./src/err/mod.rs", line 300, in pyo3::err::PyErr::take
+    let mut backtrace = backtrace::Backtrace::new();
+  File "[backtrace-0.3.76]/capture.rs", line 259, in backtrace::capture::Backtrace::new
+    let mut bt = Self::create(Self::new as usize);
+  File "[backtrace-0.3.76]/capture.rs", line 294, in backtrace::capture::Backtrace::create
+    trace(|frame| {
+  File "[backtrace-0.3.76]/backtrace/mod.rs", line 53, in backtrace::backtrace::trace
+    unsafe { trace_unsynchronized(cb) }
+  File "[backtrace-0.3.76]/backtrace/mod.rs", line 66, in backtrace::backtrace::trace_unsynchronized
+    unsafe { trace_imp(&mut cb) }
+  File "[backtrace-0.3.76]/backtrace/libunwind.rs", line 117, in backtrace::backtrace::libunwind::trace
+    uw::_Unwind_Backtrace(trace_fn, addr_of_mut!(cb).cast());
+  File "<string>", line 4, in <module>
+  File "<string>", line 2, in python_func
+  File "./tests/test_backtrace.rs", line 11, in test_backtrace::test_new_err_function_err_result::__pyfunction_produce_err_result
+    #[pyfunction]
+  File "./tests/test_backtrace.rs", line 13, in test_backtrace::test_new_err_function_err_result::produce_err_result
+    Err(PyValueError::new_err("Error result"))
+  File "./src/exceptions.rs", line 31, in pyo3::exceptions::PyValueError::new_err
+    $crate::PyErr::new::<$name, A>(args)
+  File "./src/err/mod.rs", line 131, in pyo3::err::PyErr::new
+    let backtrace = backtrace::Backtrace::new_unresolved();
+  File "[backtrace-0.3.76]/capture.rs", line 289, in backtrace::capture::Backtrace::new_unresolved
+    Self::create(Self::new_unresolved as usize)
+  File "[backtrace-0.3.76]/capture.rs", line 294, in backtrace::capture::Backtrace::create
+    trace(|frame| {
+  File "[backtrace-0.3.76]/backtrace/mod.rs", line 53, in backtrace::backtrace::trace
+    unsafe { trace_unsynchronized(cb) }
+  File "[backtrace-0.3.76]/backtrace/mod.rs", line 66, in backtrace::backtrace::trace_unsynchronized
+    unsafe { trace_imp(&mut cb) }
+  File "[backtrace-0.3.76]/backtrace/libunwind.rs", line 117, in backtrace::backtrace::libunwind::trace
+    uw::_Unwind_Backtrace(trace_fn, addr_of_mut!(cb).cast());

tests/test_backtrace.rs

@@ -0,0 +1,43 @@
+#![cfg(feature = "macros")]
+use insta::assert_snapshot;
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+
+#[test]
+fn test_new_err_function_err_result() {
+    use pyo3::prelude::PyDictMethods;
+    use pyo3::{pyfunction, types::PyDict, Python};
+
+    #[pyfunction]
+    fn produce_err_result() -> PyResult<()> {
+        Err(PyValueError::new_err("Error result"))
+    }
+
+    Python::attach(|py| {
+        let func = wrap_pyfunction!(produce_err_result)(py).unwrap();
+        let globals = PyDict::new(py);
+        globals.set_item("func", func).unwrap();
+
+        let err = py
+            .run(
+                c"def python_func():\n  func()\n\npython_func()",
+                Some(&globals),
+                None,
+            )
+            .unwrap_err();
+
+        let traceback = err.traceback(py).unwrap().format().unwrap();
+
+        insta::with_settings!({
+            snapshot_suffix => std::env::consts::OS,
+            filters => [
+                (&*std::env::current_dir().unwrap().to_string_lossy(), "."),
+                ("(?:/[\\w\\-\\.]*)+/library/core/src", "[RUST_CORE]"),
+                ("(?:/[\\w\\-\\.]*)+/library/alloc/src", "[RUST_ALLOC]"),
+                ("(?:/[\\w\\-\\.]*)+/registry/src/index\\.crates\\.io-\\w{16}/([\\w\\-\\.]*)/src", "[$1]"),
+            ],
+        }, {
+            assert_snapshot!(traceback);
+        });
+    });
+}