Update traceback.tb_next to be writable.

tgsong827 · tgsong827 · commit 94f1d7d6d83e · 2022-08-17T23:27:50.000+09:00
diff --git a/Lib/unittest/result.py b/Lib/unittest/result.py
@@ -233,9 +233,8 @@ def _remove_unittest_tb_frames(self, tb):
         while tb and not self._is_relevant_tb_level(tb):
             prev = tb
             tb = tb.tb_next
-        # TODO: RUSTPYTHON; traceback.tb_next is not writable yet #3857
-        # if prev is not None:
-        #     prev.tb_next = None
+        if prev is not None:
+            prev.tb_next = None
 
     def __repr__(self):
         return ("<%s run=%i errors=%i failures=%i>" %
diff --git a/Lib/unittest/test/test_result.py b/Lib/unittest/test/test_result.py
@@ -220,8 +220,6 @@ def test_1(self):
         self.assertIs(test_case, test)
         self.assertIsInstance(formatted_exc, str)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_addFailure_filter_traceback_frames(self):
         class Foo(unittest.TestCase):
             def test_1(self):
@@ -248,8 +246,6 @@ def get_exc_info():
         self.assertEqual(len(dropped), 1)
         self.assertIn("raise self.failureException(msg)", dropped[0])
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_addFailure_filter_traceback_frames_context(self):
         class Foo(unittest.TestCase):
             def test_1(self):
diff --git a/vm/src/builtins/traceback.rs b/vm/src/builtins/traceback.rs
@@ -1,10 +1,12 @@
+use rustpython_common::lock::PyMutex;
+
 use super::PyType;
 use crate::{class::PyClassImpl, frame::FrameRef, Context, Py, PyPayload, PyRef, VirtualMachine};
 
 #[pyclass(module = false, name = "traceback")]
 #[derive(Debug)]
 pub struct PyTraceback {
-    pub next: Option<PyTracebackRef>, // TODO: Make mutable
+    pub next: PyMutex<Option<PyTracebackRef>>,
     pub frame: FrameRef,
     pub lasti: u32,
     pub lineno: usize,
@@ -22,7 +24,7 @@ impl PyPayload for PyTraceback {
 impl PyTraceback {
     pub fn new(next: Option<PyRef<Self>>, frame: FrameRef, lasti: u32, lineno: usize) -> Self {
         PyTraceback {
-            next,
+            next: PyMutex::new(next),
             frame,
             lasti,
             lineno,
@@ -46,13 +48,18 @@ impl PyTraceback {
 
     #[pyproperty]
     fn tb_next(&self) -> Option<PyRef<Self>> {
-        self.next.as_ref().cloned()
+        self.next.lock().as_ref().cloned()
+    }
+
+    #[pyproperty(setter)]
+    fn set_tb_next(&self, value: Option<PyRef<Self>>) {
+        *self.next.lock() = value;
     }
 }
 
 impl PyTracebackRef {
     pub fn iter(&self) -> impl Iterator<Item = PyTracebackRef> {
-        std::iter::successors(Some(self.clone()), |tb| tb.next.clone())
+        std::iter::successors(Some(self.clone()), |tb| tb.next.lock().clone())
     }
 }
 
diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -161,7 +161,7 @@ fn remove_importlib_frames_inner(
     let file_name = traceback.frame.code.source_path.as_str();
 
     let (inner_tb, mut now_in_importlib) =
-        remove_importlib_frames_inner(vm, traceback.next.clone(), always_trim);
+        remove_importlib_frames_inner(vm, traceback.next.lock().clone(), always_trim);
     if file_name == "_frozen_importlib" || file_name == "_frozen_importlib_external" {
         if traceback.frame.code.obj_name.as_str() == "_call_with_frames_removed" {
             now_in_importlib = true;
diff --git a/vm/src/suggestion.rs b/vm/src/suggestion.rs
@@ -53,7 +53,7 @@ pub fn offer_suggestions(exc: &PyBaseExceptionRef, vm: &VirtualMachine) -> Optio
     } else if exc.class().is(vm.ctx.exceptions.name_error) {
         let name = exc.as_object().to_owned().get_attr("name", vm).unwrap();
         let mut tb = exc.traceback().unwrap();
-        while let Some(traceback) = tb.next.clone() {
+        for traceback in tb.iter() {
             tb = traceback;
         }
 

Original file line number	Diff line number	Diff line change
`@@ -53,7 +53,7 @@ pub fn offer_suggestions(exc: &PyBaseExceptionRef, vm: &VirtualMachine) -> Optio`
`53`	`53`	`} else if exc.class().is(vm.ctx.exceptions.name_error) {`
`54`	`54`	`let name = exc.as_object().to_owned().get_attr("name", vm).unwrap();`
`55`	`55`	`let mut tb = exc.traceback().unwrap();`
`56`		`- while let Some(traceback) = tb.next.clone() {`
	`56`	`+ for traceback in tb.iter() {`
`57`	`57`	`tb = traceback;`
`58`	`58`	`}`
`59`	`59`