Merge pull request #1977 from davidhewitt/traceback-type

types: add PyTraceback

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Add implementations for `Py::as_ref()` and `Py::into_ref()` for `Py<PySequence>`, `Py<PyIterator>` and `Py<PyMapping>`. [#1682](https://github.com/PyO3/pyo3/pull/1682)
+- Add `PyTraceback` type to represent and format Python tracebacks. [#1977](https://github.com/PyO3/pyo3/pull/1977)
 
 ### Changed
 

src/err/err_state.rs

@@ -2,15 +2,15 @@ use crate::{
     exceptions::{PyBaseException, PyTypeError},
     ffi,
     type_object::PyTypeObject,
-    types::PyType,
+    types::{PyTraceback, PyType},
     AsPyPointer, IntoPy, IntoPyPointer, Py, PyObject, Python,
 };
 
 #[derive(Clone)]
 pub(crate) struct PyErrStateNormalized {
     pub ptype: Py<PyType>,
     pub pvalue: Py<PyBaseException>,
-    pub ptraceback: Option<PyObject>,
+    pub ptraceback: Option<Py<PyTraceback>>,
 }
 
 pub(crate) enum PyErrState {

src/err/mod.rs

@@ -2,7 +2,7 @@
 
 use crate::panic::PanicException;
 use crate::type_object::PyTypeObject;
-use crate::types::PyType;
+use crate::types::{PyTraceback, PyType};
 use crate::{
     exceptions::{self, PyBaseException},
     ffi,
@@ -201,7 +201,7 @@ impl PyErr {
     ///     assert_eq!(err.ptraceback(py), None);
     /// });
     /// ```
-    pub fn ptraceback<'py>(&'py self, py: Python<'py>) -> Option<&'py PyAny> {
+    pub fn ptraceback<'py>(&'py self, py: Python<'py>) -> Option<&'py PyTraceback> {
         self.normalized(py)
             .ptraceback
             .as_ref()
@@ -497,7 +497,7 @@ impl PyErr {
             *self_state = Some(PyErrState::Normalized(PyErrStateNormalized {
                 ptype: Py::from_owned_ptr_or_opt(py, ptype).expect("Exception type missing"),
                 pvalue: Py::from_owned_ptr_or_opt(py, pvalue).expect("Exception value missing"),
-                ptraceback: PyObject::from_owned_ptr_or_opt(py, ptraceback),
+                ptraceback: Py::from_owned_ptr_or_opt(py, ptraceback),
             }));
 
             match self_state {

src/types/mod.rs

@@ -26,6 +26,7 @@ pub use self::slice::{PySlice, PySliceIndices};
 #[cfg(all(not(Py_LIMITED_API), target_endian = "little"))]
 pub use self::string::PyStringData;
 pub use self::string::{PyString, PyString as PyUnicode};
+pub use self::traceback::PyTraceback;
 pub use self::tuple::PyTuple;
 pub use self::typeobject::PyType;
 
@@ -237,5 +238,6 @@ mod sequence;
 mod set;
 mod slice;
 mod string;
+mod traceback;
 mod tuple;
 mod typeobject;

src/types/traceback.rs

@@ -0,0 +1,82 @@
+// Copyright (c) 2017-present PyO3 Project and Contributors
+
+use crate::err::{error_on_minusone, PyResult};
+use crate::ffi;
+use crate::types::PyString;
+use crate::{AsPyPointer, PyAny};
+
+/// Represents a Python traceback.
+#[repr(transparent)]
+pub struct PyTraceback(PyAny);
+
+pyobject_native_type_core!(
+    PyTraceback,
+    ffi::PyTraceBack_Type,
+    #checkfunction=ffi::PyTraceBack_Check
+);
+
+impl PyTraceback {
+    /// Formats the traceback as a string.
+    ///
+    /// This does not include the exception type and value. The exception type and value can be
+    /// formatted using the `Display` implementation for `PyErr`.
+    ///
+    /// # Example
+    ///
+    /// The following code formats a Python traceback and exception pair from Rust:
+    ///
+    /// ```rust
+    /// # use pyo3::{Python, PyResult};
+    /// # let result: PyResult<()> =
+    /// Python::with_gil(|py| {
+    ///     let err = py
+    ///         .run("raise Exception('banana')", None, None)
+    ///         .expect_err("raise will create a Python error");
+    ///
+    ///     let traceback = err.ptraceback(py).expect("raised exception will have a traceback");
+    ///     assert_eq!(
+    ///         format!("{}{}", traceback.format()?, err),
+    ///         "\
+    /// Traceback (most recent call last):
+    ///   File \"<string>\", line 1, in <module>
+    /// Exception: banana\
+    /// "
+    ///     );
+    ///     Ok(())
+    /// })
+    /// # ;
+    /// # result.expect("example failed");
+    /// ```
+    pub fn format(&self) -> PyResult<String> {
+        let py = self.py();
+        let string_io = py.import("io")?.getattr("StringIO")?.call0()?;
+        let result = unsafe { ffi::PyTraceBack_Print(self.as_ptr(), string_io.as_ptr()) };
+        error_on_minusone(py, result)?;
+        let formatted = string_io
+            .getattr("getvalue")?
+            .call0()?
+            .downcast::<PyString>()?
+            .to_str()?
+            .to_owned();
+        Ok(formatted)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::Python;
+
+    #[test]
+    fn format_traceback() {
+        Python::with_gil(|py| {
+            let err = py
+                .run("raise Exception('banana')", None, None)
+                .expect_err("raising should have given us an error");
+
+            assert_eq!(
+                err.ptraceback(py).unwrap().format().unwrap(),
+                "Traceback (most recent call last):\n  File \"<string>\", line 1, in <module>\n"
+            );
+        })
+    }
+}