From f3e4b4f4fce0922a07b72ee9f0ddc5c15846c433 Mon Sep 17 00:00:00 2001
From: Jeong YunWon <jeong@youknowone.org>
Date: Tue, 24 Jun 2025 15:09:13 +0900
Subject: [PATCH] Exception.set_traceback_typed

---
 Lib/test/test_exceptions.py |  2 --
 vm/src/exceptions.rs        | 25 ++++++++++++++++++++++---
 vm/src/frame.rs             |  2 +-
 vm/src/import.rs            |  2 +-
 4 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/Lib/test/test_exceptions.py b/Lib/test/test_exceptions.py
index eb377877bbc..1864e611ab0 100644
--- 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
index 885f82b71ad..1a6c8ed7ee7 100644
--- 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
index 925899564a0..2ffd0335464 100644
--- 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
index 90aadbdbf27..69c8396c07a 100644
--- 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);
     }
 }