Change remove_importlib_frames to work with tb object

vm/src/frame.rs

@@ -14,7 +14,7 @@ use crate::obj::objlist;
 use crate::obj::objslice::PySlice;
 use crate::obj::objstr;
 use crate::obj::objstr::PyString;
-use crate::obj::objtraceback::PyTraceback;
+use crate::obj::objtraceback::{PyTraceback, PyTracebackRef};
 use crate::obj::objtuple::PyTuple;
 use crate::obj::objtype;
 use crate::obj::objtype::PyClassRef;
@@ -191,12 +191,21 @@ impl Frame {
                     let traceback = vm
                         .get_attribute(exception.clone(), "__traceback__")
                         .unwrap();
+
+                    let next = if vm.is_none(&traceback) {
+                        None
+                    } else {
+                        let traceback: PyTracebackRef = traceback
+                            .downcast()
+                            .expect("next must be a traceback object");
+                        Some(traceback)
+                    };
+
                     let new_traceback = PyTraceback::new(
-                        traceback,
+                        next,
                         self.clone().into_ref(vm),
                         self.lasti.get(),
                         lineno.row(),
-                        vm,
                     );
                     vm.set_attr(&exception, "__traceback__", new_traceback.into_ref(vm))
                         .unwrap();

vm/src/import.rs

@@ -4,7 +4,8 @@
 use rand::Rng;
 
 use crate::bytecode::CodeObject;
-use crate::obj::{objcode, objsequence, objstr, objtype};
+use crate::obj::objtraceback::{PyTraceback, PyTracebackRef};
+use crate::obj::{objcode, objtype};
 use crate::pyobject::{ItemProtocol, PyObjectRef, PyResult, PyValue};
 use crate::scope::Scope;
 use crate::version::get_git_revision;
@@ -100,36 +101,50 @@ pub fn import_codeobj(
     Ok(module)
 }
 
+fn remove_importlib_frames_inner(
+    vm: &VirtualMachine,
+    tb: Option<PyTracebackRef>,
+    always_trim: bool,
+) -> Option<PyTracebackRef> {
+    if tb.is_none() {
+        return None;
+    }
+    let traceback = tb.unwrap();
+    let file_name = traceback.frame.code.source_path.to_string();
+    if file_name == "_frozen_importlib" || file_name == "_frozen_importlib_external" {
+        if always_trim || traceback.frame.code.obj_name == "_call_with_frames_removed" {
+            return remove_importlib_frames_inner(
+                vm,
+                traceback.next.as_ref().map(|x| x.clone()),
+                always_trim,
+            );
+        }
+    }
+
+    Some(
+        PyTraceback::new(
+            remove_importlib_frames_inner(
+                vm,
+                traceback.next.as_ref().map(|x| x.clone()),
+                always_trim,
+            ),
+            traceback.frame.clone(),
+            traceback.lasti,
+            traceback.lineno,
+        )
+        .into_ref(vm),
+    )
+}
+
 // TODO: This function should do nothing on verbose mode.
 pub fn remove_importlib_frames(vm: &VirtualMachine, exc: &PyObjectRef) -> PyObjectRef {
     let always_trim = objtype::isinstance(exc, &vm.ctx.exceptions.import_error);
 
     if let Ok(tb) = vm.get_attribute(exc.clone(), "__traceback__") {
-        if objtype::isinstance(&tb, &vm.ctx.list_type()) {
-            let tb_entries = objsequence::get_elements_list(&tb).to_vec();
-            let mut in_importlib = false;
-            let new_tb = tb_entries
-                .iter()
-                .filter(|tb_entry| {
-                    let location_attrs = objsequence::get_elements_tuple(&tb_entry);
-                    let file_name = objstr::get_value(&location_attrs[0]);
-                    if file_name == "_frozen_importlib" || file_name == "_frozen_importlib_external"
-                    {
-                        let run_obj_name = objstr::get_value(&location_attrs[2]);
-                        if run_obj_name == "_call_with_frames_removed" {
-                            in_importlib = true;
-                        }
-                        !always_trim && !in_importlib
-                    } else {
-                        in_importlib = false;
-                        true
-                    }
-                })
-                .cloned()
-                .collect();
-            vm.set_attr(exc, "__traceback__", vm.ctx.new_list(new_tb))
-                .unwrap();
-        }
+        let base_tb: PyTracebackRef = tb.downcast().expect("must be a traceback object");
+        let trimed_tb = remove_importlib_frames_inner(vm, Some(base_tb), always_trim)
+            .map_or(vm.get_none(), |x| x.into_object());
+        vm.set_attr(exc, "__traceback__", trimed_tb).unwrap();
     }
     exc.clone()
 }

vm/src/obj/objtraceback.rs

@@ -1,6 +1,6 @@
 use crate::frame::FrameRef;
 use crate::obj::objtype::PyClassRef;
-use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyValue};
+use crate::pyobject::{PyClassImpl, PyContext, PyRef, PyValue};
 use crate::vm::VirtualMachine;
 
 #[pyclass]
@@ -22,21 +22,7 @@ impl PyValue for PyTraceback {
 
 #[pyimpl]
 impl PyTraceback {
-    pub fn new(
-        next: PyObjectRef,
-        frame: FrameRef,
-        lasti: usize,
-        lineno: usize,
-        vm: &VirtualMachine,
-    ) -> Self {
-        let next = if vm.is_none(&next) {
-            None
-        } else {
-            let traceback: PyTracebackRef =
-                next.downcast().expect("next must be a traceback object");
-            Some(traceback)
-        };
-
+    pub fn new(next: Option<PyTracebackRef>, frame: FrameRef, lasti: usize, lineno: usize) -> Self {
         PyTraceback {
             next,
             frame,