From 94f1d7d6d83e1a7b8697dc19756c77d2082c1faa Mon Sep 17 00:00:00 2001
From: TaeGit <tgsong827@gmail.com>
Date: Wed, 17 Aug 2022 23:01:57 +0900
Subject: [PATCH] Update traceback.tb_next to be writable.

---
 Lib/unittest/result.py           |  5 ++---
 Lib/unittest/test/test_result.py |  4 ----
 vm/src/builtins/traceback.rs     | 15 +++++++++++----
 vm/src/import.rs                 |  2 +-
 vm/src/suggestion.rs             |  2 +-
 5 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/Lib/unittest/result.py b/Lib/unittest/result.py
index 3b47c27f93..3da7005e60 100644
--- 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
index d3b56c7e7e..c5aaba0ff5 100644
--- 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
index 98d9a3aedb..f4e21e7191 100644
--- 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
index e7c29d2f70..3068cd275c 100644
--- 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
index b23f9c447b..34aeddd0f8 100644
--- 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;
         }