Exception.set_traceback_typed (#5832)

youknowone · web-flow · commit 3734f32e4226 · 2025-06-24T18:56:10.000+09:00
diff --git a/Lib/test/test_exceptions.py b/Lib/test/test_exceptions.py
@@ -607,8 +607,6 @@ class MyException(Exception):
         self.assertIsInstance(e, MyException)
         self.assertEqual(e.__traceback__, tb)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def testInvalidTraceback(self):
         try:
             Exception().__traceback__ = 5
diff --git a/vm/src/exceptions.rs b/vm/src/exceptions.rs
@@ -4,7 +4,8 @@ use crate::object::{Traverse, TraverseFn};
 use crate::{
     AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine,
     builtins::{
-        PyNone, PyStr, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef, traceback::PyTracebackRef,
+        PyNone, PyStr, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef,
+        traceback::{PyTraceback, PyTracebackRef},
     },
     class::{PyClassImpl, StaticType},
     convert::{ToPyException, ToPyObject},
@@ -324,7 +325,7 @@ impl VirtualMachine {
         let ctor = ExceptionCtor::try_from_object(self, exc_type)?;
         let exc = ctor.instantiate_value(exc_val, self)?;
         if let Some(tb) = Option::<PyTracebackRef>::try_from_object(self, exc_tb)? {
-            exc.set_traceback(Some(tb));
+            exc.set_traceback_typed(Some(tb));
         }
         Ok(exc)
     }
@@ -584,7 +585,25 @@ impl PyBaseException {
     }
 
     #[pygetset(magic, setter)]
-    pub fn set_traceback(&self, traceback: Option<PyTracebackRef>) {
+    pub fn set_traceback(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
+        let traceback = if vm.is_none(&value) {
+            None
+        } else {
+            match value.downcast::<PyTraceback>() {
+                Ok(tb) => Some(tb),
+                Err(_) => {
+                    return Err(
+                        vm.new_type_error("__traceback__ must be a traceback or None".to_owned())
+                    );
+                }
+            }
+        };
+        self.set_traceback_typed(traceback);
+        Ok(())
+    }
+
+    // Helper method for internal use that doesn't require PyObjectRef
+    pub(crate) fn set_traceback_typed(&self, traceback: Option<PyTracebackRef>) {
         *self.traceback.write() = traceback;
     }
 
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -387,7 +387,7 @@ impl ExecutingFrame<'_> {
                         let new_traceback =
                             PyTraceback::new(next, frame.object.to_owned(), frame.lasti(), loc.row);
                         vm_trace!("Adding to traceback: {:?} {:?}", new_traceback, loc.row);
-                        exception.set_traceback(Some(new_traceback.into_ref(&vm.ctx)));
+                        exception.set_traceback_typed(Some(new_traceback.into_ref(&vm.ctx)));
 
                         vm.contextualize_exception(&exception);
 
diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -213,6 +213,6 @@ pub fn remove_importlib_frames(vm: &VirtualMachine, exc: &PyBaseExceptionRef) {
 
     if let Some(tb) = exc.traceback() {
         let trimmed_tb = remove_importlib_frames_inner(vm, Some(tb), always_trim).0;
-        exc.set_traceback(trimmed_tb);
+        exc.set_traceback_typed(trimmed_tb);
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -213,6 +213,6 @@ pub fn remove_importlib_frames(vm: &VirtualMachine, exc: &PyBaseExceptionRef) {`
`213`	`213`
`214`	`214`	`if let Some(tb) = exc.traceback() {`
`215`	`215`	`let trimmed_tb = remove_importlib_frames_inner(vm, Some(tb), always_trim).0;`
`216`		`- exc.set_traceback(trimmed_tb);`
	`216`	`+ exc.set_traceback_typed(trimmed_tb);`
`217`	`217`	`}`
`218`	`218`	`}`