Add PyTraceback::append

newsfragments/5857.added.md

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

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,39 @@ 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>(),
+            )
+        }
+    }
+}
+
+#[doc(alias = "PyFrame")]
+pub trait PyFrameMethods<'py>: Sealed {
+    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()) }
+    }
+}

src/types/traceback.rs

@@ -1,6 +1,11 @@
 use crate::err::{error_on_minusone, PyResult};
 use crate::types::{any::PyAnyMethods, string::PyStringMethods, PyString};
 use crate::{ffi, Bound, PyAny};
+#[cfg(all(not(Py_LIMITED_API), not(PyPy), not(GraalPy)))]
+use crate::{
+    types::{frame::PyFrameMethods, PyFrame},
+    BoundObject, IntoPyObject, PyTypeCheck, Python,
+};
 
 /// Represents a Python traceback.
 ///
@@ -20,6 +25,43 @@ pyobject_native_type_core!(
     #checkfunction=ffi::PyTraceBack_Check
 );
 
+impl PyTraceback {
+    #[cfg(all(not(Py_LIMITED_API), not(PyPy), not(GraalPy)))]
+    pub(crate) fn new<'py>(
+        py: Python<'py>,
+        next: Option<Bound<'py, PyTraceback>>,
+        frame: Bound<'py, PyFrame>,
+        instruction_index: i32,
+        line_number: i32,
+    ) -> PyResult<Bound<'py, PyTraceback>> {
+        unsafe {
+            Ok(PyTraceback::classinfo_object(py)
+                .call1((next, frame, instruction_index, line_number))?
+                .cast_into_unchecked())
+        }
+    }
+
+    /// Creates a new traceback object from an iterator of frames.
+    ///
+    /// The frames should be ordered from newest to oldest, i.e. the first frame in the iterator
+    /// will be the innermost frame in the traceback.
+    #[cfg(all(not(Py_LIMITED_API), not(PyPy), not(GraalPy)))]
+    pub fn from_frames<'py, I>(
+        py: Python<'py>,
+        frames: I,
+    ) -> PyResult<Option<Bound<'py, PyTraceback>>>
+    where
+        I: IntoIterator,
+        I::Item: IntoPyObject<'py, Target = PyFrame>,
+    {
+        frames.into_iter().try_fold(None, |prev, frame| {
+            let frame = frame.into_pyobject(py).map_err(Into::into)?.into_bound();
+            let line_number = frame.line_number();
+            PyTraceback::new(py, prev, frame, 0, line_number).map(Some)
+        })
+    }
+}
+
 /// Implementation of functionality for [`PyTraceback`].
 ///
 /// These methods are defined for the `Bound<'py, PyTraceback>` smart pointer, so to use method call
@@ -59,6 +101,16 @@ pub trait PyTracebackMethods<'py>: crate::sealed::Sealed {
     /// # result.expect("example failed");
     /// ```
     fn format(&self) -> PyResult<String>;
+
+    /// Get the next traceback towards the frame where the exception was raised.
+    fn next_traceback(&self) -> PyResult<Option<Bound<'py, PyTraceback>>>;
+
+    /// Get the innermost traceback, i.e. the traceback corresponding to the frame where the exception was raised.
+    fn innermost_traceback(&self) -> PyResult<Bound<'py, PyTraceback>>;
+
+    /// Append a traceback to the end of this traceback, i.e. the frames of the appended traceback
+    /// will be newer than the frames of this traceback.
+    fn append(&self, traceback: Bound<'py, PyTraceback>) -> PyResult<()>;
 }
 
 impl<'py> PyTracebackMethods<'py> for Bound<'py, PyTraceback> {
@@ -78,13 +130,31 @@ impl<'py> PyTracebackMethods<'py> for Bound<'py, PyTraceback> {
             .into_owned();
         Ok(formatted)
     }
+
+    fn next_traceback(&self) -> PyResult<Option<Bound<'py, PyTraceback>>> {
+        Ok(self.getattr(intern!(self.py(), "tb_next"))?.extract()?)
+    }
+
+    fn innermost_traceback(&self) -> PyResult<Bound<'py, PyTraceback>> {
+        let mut current = self.clone();
+        while let Some(next) = current.next_traceback()? {
+            current = next;
+        }
+        Ok(current)
+    }
+
+    fn append(&self, next: Bound<'py, PyTraceback>) -> PyResult<()> {
+        self.innermost_traceback()?
+            .setattr(intern!(self.py(), "tb_next"), next)
+    }
 }
 
 #[cfg(test)]
 mod tests {
+    use super::*;
     use crate::IntoPyObject;
     use crate::{
-        types::{any::PyAnyMethods, dict::PyDictMethods, traceback::PyTracebackMethods, PyDict},
+        types::{dict::PyDictMethods, PyDict},
         PyErr, Python,
     };
 
@@ -146,4 +216,61 @@ def f():
             assert!(err_object.getattr("__traceback__").unwrap().is(&traceback));
         })
     }
+
+    #[test]
+    #[cfg(all(not(Py_LIMITED_API), not(PyPy), not(GraalPy)))]
+    fn test_create_traceback() {
+        Python::attach(|py| {
+            // most recent frame first, oldest frame last
+            let frames = [
+                PyFrame::new(py, c"file3.py", c"func3", 30).unwrap(),
+                PyFrame::new(py, c"file2.py", c"func2", 20).unwrap(),
+                PyFrame::new(py, c"file1.py", c"func1", 10).unwrap(),
+            ];
+
+            let traceback = PyTraceback::from_frames(py, frames).unwrap().unwrap();
+            assert_eq!(
+                traceback.format().unwrap(), "Traceback (most recent call last):\n  File \"file1.py\", line 10, in func1\n  File \"file2.py\", line 20, in func2\n  File \"file3.py\", line 30, in func3\n"
+            );
+        })
+    }
+
+    #[test]
+    #[cfg(all(not(Py_LIMITED_API), not(PyPy), not(GraalPy)))]
+    fn test_insert_traceback() {
+        Python::attach(|py| {
+            let traceback = PyTraceback::from_frames(
+                py,
+                [
+                    PyFrame::new(py, c"file2.py", c"func2", 20).unwrap(),
+                    PyFrame::new(py, c"file1.py", c"func1", 10).unwrap(),
+                ],
+            )
+            .unwrap()
+            .unwrap();
+
+            let rust_traceback = PyTraceback::from_frames(
+                py,
+                [
+                    PyFrame::new(py, c"rust2.rs", c"func2", 22).unwrap(),
+                    PyFrame::new(py, c"rust1.rs", c"func1", 11).unwrap(),
+                ],
+            )
+            .unwrap()
+            .unwrap();
+
+            // Stacktrace where python calls into rust
+            traceback.append(rust_traceback).unwrap();
+
+            assert_eq!(
+                traceback.format().unwrap(),
+                r#"Traceback (most recent call last):
+  File "file1.py", line 10, in func1
+  File "file2.py", line 20, in func2
+  File "rust1.rs", line 11, in func1
+  File "rust2.rs", line 22, in func2
+"#
+            );
+        })
+    }
 }