Inject rust stack frames into PyErr traceback

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"] }

newsfragments/5857.added.md

@@ -0,0 +1 @@
+Allow constructing `PyTraceback`

src/err/err_state.rs

@@ -4,6 +4,8 @@ use std::{
     thread::ThreadId,
 };
 
+#[cfg(all(debug_assertions, not(Py_LIMITED_API)))]
+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(all(debug_assertions, not(Py_LIMITED_API)))]
+        let backtrace = backtrace::Backtrace::new_unresolved();
+
         Self::from_inner(PyErrStateInner::Lazy(Box::new(move |py| {
             PyErrStateLazyFnOutput {
                 ptype,
                 pvalue: args.arguments(py),
+                #[cfg(all(debug_assertions, not(Py_LIMITED_API)))]
+                backtrace,
             }
         })))
     }
@@ -301,6 +308,8 @@ impl PyErrStateNormalized {
 pub(crate) struct PyErrStateLazyFnOutput {
     pub(crate) ptype: Py<PyAny>,
     pub(crate) pvalue: Py<PyAny>,
+    #[cfg(all(debug_assertions, not(Py_LIMITED_API)))]
+    pub(crate) backtrace: backtrace::Backtrace,
 }
 
 pub(crate) type PyErrStateLazyFn =
@@ -390,15 +399,62 @@ 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(all(debug_assertions, not(Py_LIMITED_API)))]
+        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())
+            #[cfg(all(debug_assertions, not(Py_LIMITED_API)))]
+            let traceback =
+                PyTraceback::from_frames(py, None, backtrace_to_frames(py, &mut backtrace))
+                    .ok()
+                    .flatten()
+                    .map_or_else(std::ptr::null_mut, Bound::into_ptr);
+
+            #[cfg(not(Py_3_12))]
+            {
+                #[cfg(all(debug_assertions, not(Py_LIMITED_API)))]
+                ffi::PyErr_Restore(ptype.into_ptr(), pvalue.into_ptr(), traceback);
+
+                #[cfg(not(all(debug_assertions, not(Py_LIMITED_API))))]
+                ffi::PyErr_SetObject(ptype.as_ptr(), pvalue.as_ptr());
+            }
+
+            #[cfg(Py_3_12)]
+            {
+                let exc = if ffi::PyExceptionInstance_Check(pvalue.as_ptr()) != 0 {
+                    // If it's already an exception instance, keep it as-is.
+                    ffi::Py_NewRef(pvalue.as_ptr())
+                } else if pvalue.as_ptr() == ffi::Py_None() {
+                    // If the value is None, call the type with no arguments.
+                    ffi::PyObject_CallNoArgs(ptype.as_ptr())
+                } else if ffi::PyTuple_Check(pvalue.as_ptr()) != 0 {
+                    // If the value is a tuple, unpack it as arguments to the type.
+                    ffi::PyObject_Call(ptype.as_ptr(), pvalue.as_ptr(), std::ptr::null_mut())
+                } else {
+                    // Fallback: type(value)
+                    ffi::PyObject_CallOneArg(ptype.as_ptr(), pvalue.as_ptr())
+                };
+
+                if exc.is_null() {
+                    // Exception constructor raised an exception, so propagate that instead of the original one.
+                    return;
+                }
+
+                #[cfg(all(debug_assertions, not(Py_LIMITED_API)))]
+                ffi::PyException_SetTraceback(exc, traceback);
+
+                ffi::PyErr_SetRaisedException(exc);
+            }
         }
     }
 }

src/err/mod.rs

@@ -20,6 +20,8 @@ use crate::{BoundObject, Py, PyAny, Python};
 use err_state::{PyErrState, PyErrStateLazyFnOutput, PyErrStateNormalized};
 use std::convert::Infallible;
 use std::ffi::CStr;
+#[cfg(all(debug_assertions, not(Py_LIMITED_API)))]
+use {crate::types::PyFrame, std::ffi::CString};
 
 mod cast_error;
 mod downcast_error;
@@ -127,10 +129,14 @@ impl PyErr {
         T: PyTypeInfo,
         A: PyErrArguments + Send + Sync + 'static,
     {
+        #[cfg(all(debug_assertions, not(Py_LIMITED_API)))]
+        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(all(debug_assertions, not(Py_LIMITED_API)))]
+                backtrace,
             }
         })))
     }
@@ -289,7 +295,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(all(debug_assertions, not(Py_LIMITED_API)))]
+        {
+            let mut backtrace = backtrace::Backtrace::new();
+            if let Some(traceback) = PyTraceback::from_frames(
+                py,
+                err.traceback(py),
+                backtrace_to_frames(py, &mut backtrace),
+            )
+            .ok()
+            .flatten()
+            {
+                err.set_traceback(py, Some(traceback));
+            }
+        }
+
+        Some(err)
     }
 
     #[cold]
@@ -696,6 +719,49 @@ impl<'py> IntoPyObject<'py> for PyErr {
     }
 }
 
+#[cfg(all(debug_assertions, not(Py_LIMITED_API)))]
+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())
+        .map(|symbol| (symbol.name().map(|name| format!("{name:#}")), symbol))
+        .skip_while(|(name, _)| {
+            if cfg!(any(target_vendor = "apple", windows)) {
+                // On Apple & Windows platforms, backtrace is not able to remove internal frames
+                // from the backtrace, so we need to skip them manually here.
+                name.as_ref()
+                    .map(|name| name.starts_with("backtrace::"))
+                    .unwrap_or(true)
+            } else {
+                false
+            }
+        })
+        // The first frame is always the capture function, so skip it.
+        .skip(1)
+        .take_while(|(name, _)| {
+            name.as_ref()
+                .map(|name| {
+                    !(name.starts_with("pyo3::impl_::trampoline::")
+                        || name.contains("__rust_begin_short_backtrace"))
+                })
+                .unwrap_or(true)
+        })
+        .filter_map(move |(name, symbol)| {
+            let file =
+                CString::new(symbol.filename()?.as_os_str().to_string_lossy().as_ref()).ok()?;
+
+            let function = CString::new(name.as_deref().unwrap_or("<unknown>")).ok()?;
+            let line = symbol.lineno()?;
+
+            PyFrame::new(py, &file, &function, line as _).ok()
+        })
+}
+
 impl<'py> IntoPyObject<'py> for &PyErr {
     type Target = PyBaseException;
     type Output = Bound<'py, Self::Target>;
@@ -844,6 +910,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg(false)]
     fn err_debug() {
         // Debug representation should be like the following (without the newlines):
         // PyErr {

src/sealed.rs

@@ -1,4 +1,6 @@
 use crate::impl_::pyfunction::PyFunctionDef;
+#[cfg(all(not(Py_LIMITED_API), not(PyPy), not(GraalPy)))]
+use crate::types::PyFrame;
 use crate::types::{
     PyBool, PyByteArray, PyBytes, PyCapsule, PyComplex, PyDict, PyFloat, PyFrozenSet, PyList,
     PyMapping, PyMappingProxy, PyModule, PyRange, PySequence, PySet, PySlice, PyString,
@@ -42,6 +44,8 @@ impl Sealed for Bound<'_, PySet> {}
 impl Sealed for Bound<'_, PySlice> {}
 impl Sealed for Bound<'_, PyString> {}
 impl Sealed for Bound<'_, PyTraceback> {}
+#[cfg(all(not(Py_LIMITED_API), not(PyPy), not(GraalPy)))]
+impl Sealed for Bound<'_, PyFrame> {}
 impl Sealed for Bound<'_, PyTuple> {}
 impl Sealed for Bound<'_, PyType> {}
 impl Sealed for Bound<'_, PyWeakref> {}

src/types/frame.rs

@@ -1,5 +1,10 @@
-use crate::ffi;
+use crate::ffi_ptr_ext::FfiPtrExt;
+use crate::sealed::Sealed;
+use crate::types::PyDict;
 use crate::PyAny;
+use crate::{ffi, Bound, PyResult, Python};
+use pyo3_ffi::PyObject;
+use std::ffi::CStr;
 
 /// Represents a Python frame.
 ///
@@ -15,3 +20,58 @@ pyobject_native_type_core!(
     "FrameType",
     #checkfunction=ffi::PyFrame_Check
 );
+
+impl PyFrame {
+    /// Creates a new frame object.
+    pub fn new<'py>(
+        py: Python<'py>,
+        file_name: &CStr,
+        func_name: &CStr,
+        line_number: i32,
+    ) -> PyResult<Bound<'py, PyFrame>> {
+        // Safety: Thread is attached because we have a python token
+        let state = unsafe { ffi::compat::PyThreadState_GetUnchecked() };
+        let globals = PyDict::new(py);
+        let locals = PyDict::new(py);
+
+        unsafe {
+            let code = ffi::PyCode_NewEmpty(file_name.as_ptr(), func_name.as_ptr(), line_number);
+            Ok(
+                ffi::PyFrame_New(state, code, globals.as_ptr(), locals.as_ptr())
+                    .cast::<PyObject>()
+                    .assume_owned_or_err(py)?
+                    .cast_into_unchecked::<PyFrame>(),
+            )
+        }
+    }
+}
+
+/// Implementation of functionality for [`PyFrame`].
+///
+/// These methods are defined for the `Bound<'py, PyFrame>` smart pointer, so to use method call
+/// syntax these methods are separated into a trait, because stable Rust does not yet support
+/// `arbitrary_self_types`.
+#[doc(alias = "PyFrame")]
+pub trait PyFrameMethods<'py>: Sealed {
+    /// Returns the line number of the current instruction in the frame.
+    fn line_number(&self) -> i32;
+}
+
+impl<'py> PyFrameMethods<'py> for Bound<'py, PyFrame> {
+    fn line_number(&self) -> i32 {
+        unsafe { ffi::PyFrame_GetLineNumber(self.as_ptr().cast()) }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_frame_creation() {
+        Python::attach(|py| {
+            let frame = PyFrame::new(py, c"file.py", c"func", 42).unwrap();
+            assert_eq!(frame.line_number(), 42);
+        });
+    }
+}

src/types/mod.rs

@@ -16,7 +16,7 @@ pub use self::dict::{PyDictItems, PyDictKeys, PyDictValues};
 pub use self::ellipsis::PyEllipsis;
 pub use self::float::{PyFloat, PyFloatMethods};
 #[cfg(all(not(Py_LIMITED_API), not(PyPy), not(GraalPy)))]
-pub use self::frame::PyFrame;
+pub use self::frame::{PyFrame, PyFrameMethods};
 pub use self::frozenset::{PyFrozenSet, PyFrozenSetBuilder, PyFrozenSetMethods};
 pub use self::function::PyCFunction;
 #[cfg(not(Py_LIMITED_API))]