From 7e52b6e419cef5531ec3dcd1ad3d7c7134f02ad8 Mon Sep 17 00:00:00 2001
From: Nicolas Avrutin <nicolas@avrutin.net>
Date: Tue, 18 Feb 2025 15:30:29 -0500
Subject: [PATCH] Format python traceback in impl Debug for PyErr. (#4900)

Fixes #4863.
---
 newsfragments/4900.changed.md           |  1 +
 src/err/mod.rs                          | 35 +++++++++----
 tests/test_pyerr_debug_unformattable.rs | 67 +++++++++++++++++++++++++
 3 files changed, 94 insertions(+), 9 deletions(-)
 create mode 100644 newsfragments/4900.changed.md
 create mode 100644 tests/test_pyerr_debug_unformattable.rs

diff --git a/newsfragments/4900.changed.md b/newsfragments/4900.changed.md
new file mode 100644
index 00000000000..89bab779af1
--- /dev/null
+++ b/newsfragments/4900.changed.md
@@ -0,0 +1 @@
+Format python traceback in impl Debug for PyErr.
diff --git a/src/err/mod.rs b/src/err/mod.rs
index 20108f6e8dc..bd026bab14d 100644
--- a/src/err/mod.rs
+++ b/src/err/mod.rs
@@ -2,7 +2,10 @@ use crate::instance::Bound;
 use crate::panic::PanicException;
 use crate::type_object::PyTypeInfo;
 use crate::types::any::PyAnyMethods;
-use crate::types::{string::PyStringMethods, typeobject::PyTypeMethods, PyTraceback, PyType};
+use crate::types::{
+    string::PyStringMethods, traceback::PyTracebackMethods, typeobject::PyTypeMethods, PyTraceback,
+    PyType,
+};
 use crate::{
     exceptions::{self, PyBaseException},
     ffi,
@@ -770,7 +773,20 @@ impl std::fmt::Debug for PyErr {
             f.debug_struct("PyErr")
                 .field("type", &self.get_type(py))
                 .field("value", self.value(py))
-                .field("traceback", &self.traceback(py))
+                .field(
+                    "traceback",
+                    &self.traceback(py).map(|tb| match tb.format() {
+                        Ok(s) => s,
+                        Err(err) => {
+                            err.write_unraisable(py, Some(&tb));
+                            // It would be nice to format what we can of the
+                            // error, but we can't guarantee that the error
+                            // won't have another unformattable traceback inside
+                            // it and we want to avoid an infinite recursion.
+                            format!("<unformattable {:?}>", tb)
+                        }
+                    }),
+                )
                 .finish()
         })
     }
@@ -1058,7 +1074,7 @@ mod tests {
         // PyErr {
         //     type: <class 'Exception'>,
         //     value: Exception('banana'),
-        //     traceback: Some(<traceback object at 0x..)"
+        //     traceback:  Some(\"Traceback (most recent call last):\\n  File \\\"<string>\\\", line 1, in <module>\\n\")
         // }
 
         Python::with_gil(|py| {
@@ -1070,15 +1086,16 @@ mod tests {
             assert!(debug_str.starts_with("PyErr { "));
             assert!(debug_str.ends_with(" }"));
 
-            // strip "PyErr { " and " }"
-            let mut fields = debug_str["PyErr { ".len()..debug_str.len() - 2].split(", ");
+            // Strip "PyErr { " and " }". Split into 3 substrings to separate type,
+            // value, and traceback while not splitting the string within traceback.
+            let mut fields = debug_str["PyErr { ".len()..debug_str.len() - 2].splitn(3, ", ");
 
             assert_eq!(fields.next().unwrap(), "type: <class 'Exception'>");
             assert_eq!(fields.next().unwrap(), "value: Exception('banana')");
-
-            let traceback = fields.next().unwrap();
-            assert!(traceback.starts_with("traceback: Some(<traceback object at 0x"));
-            assert!(traceback.ends_with(">)"));
+            assert_eq!(
+                fields.next().unwrap(),
+                "traceback: Some(\"Traceback (most recent call last):\\n  File \\\"<string>\\\", line 1, in <module>\\n\")"
+            );
 
             assert!(fields.next().is_none());
         });
diff --git a/tests/test_pyerr_debug_unformattable.rs b/tests/test_pyerr_debug_unformattable.rs
new file mode 100644
index 00000000000..615868207b4
--- /dev/null
+++ b/tests/test_pyerr_debug_unformattable.rs
@@ -0,0 +1,67 @@
+use pyo3::ffi;
+use pyo3::prelude::*;
+
+// This test mucks around with sys.modules, so run it separately to prevent it
+// from potentially corrupting the state of the python interpreter used in other
+// tests.
+
+#[test]
+fn err_debug_unformattable() {
+    // Debug representation should be like the following (without the newlines):
+    // PyErr {
+    //     type: <class 'Exception'>,
+    //     value: Exception('banana'),
+    //     traceback: Some(\"<unformattable <traceback object at 0x...>>\")
+    // }
+
+    Python::with_gil(|py| {
+        // PyTracebackMethods::format uses io.StringIO. Mock it out to trigger a
+        // formatting failure:
+        // TypeError: 'Mock' object cannot be converted to 'PyString'
+        let err = py
+            .run(
+                ffi::c_str!(
+                    r#"
+import io, sys, unittest.mock
+sys.modules['orig_io'] = sys.modules['io']
+sys.modules['io'] = unittest.mock.Mock()
+raise Exception('banana')"#
+                ),
+                None,
+                None,
+            )
+            .expect_err("raising should have given us an error");
+
+        let debug_str = format!("{:?}", err);
+        assert!(debug_str.starts_with("PyErr { "));
+        assert!(debug_str.ends_with(" }"));
+
+        // Strip "PyErr { " and " }". Split into 3 substrings to separate type,
+        // value, and traceback while not splitting the string within traceback.
+        let mut fields = debug_str["PyErr { ".len()..debug_str.len() - 2].splitn(3, ", ");
+
+        assert_eq!(fields.next().unwrap(), "type: <class 'Exception'>");
+        assert_eq!(fields.next().unwrap(), "value: Exception('banana')");
+
+        let traceback = fields.next().unwrap();
+        assert!(
+            traceback.starts_with("traceback: Some(\"<unformattable <traceback object at 0x"),
+            "assertion failed, actual traceback str: {:?}",
+            traceback
+        );
+        assert!(fields.next().is_none());
+
+        py.run(
+            ffi::c_str!(
+                r#"
+import io, sys, unittest.mock
+sys.modules['io'] = sys.modules['orig_io']
+del sys.modules['orig_io']
+"#
+            ),
+            None,
+            None,
+        )
+        .unwrap();
+    });
+}