Remove importlib frames from traceback

palaviv · palaviv · commit c9f01cc490fe · 2019-07-04T19:20:34.000+03:00
diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -6,8 +6,8 @@ use std::path::PathBuf;
 
 use crate::bytecode::CodeObject;
 use crate::frame::Scope;
-use crate::obj::{objcode, objsequence, objstr};
-use crate::pyobject::{ItemProtocol, PyResult, PyValue};
+use crate::obj::{objcode, objsequence, objstr, objtype};
+use crate::pyobject::{ItemProtocol, PyObjectRef, PyResult, PyValue};
 use crate::util;
 use crate::vm::VirtualMachine;
 #[cfg(feature = "rustpython-compiler")]
@@ -144,3 +144,35 @@ fn find_source(vm: &VirtualMachine, current_path: PathBuf, name: &str) -> Result
         None => Err(format!("No module named '{}'", name)),
     }
 }
+
+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 mut tb_entries = objsequence::get_elements_list(&tb).to_vec();
+            tb_entries.reverse();
+            let mut new_tb = Vec::with_capacity(tb_entries.len());
+
+            for tb_entry in tb_entries.iter() {
+                let mut current_chunk = vec![];
+                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" {
+                    current_chunk.clear();
+                    new_tb.push(tb_entry.clone())
+                } else {
+                    current_chunk.push(tb_entry.clone());
+                    let run_obj_name = objstr::get_value(&location_attrs[0]);
+                    if run_obj_name == "_call_with_frames_removed" || always_trim {
+                        new_tb.append(&mut current_chunk);
+                    }
+                };
+            }
+            new_tb.reverse();
+            vm.set_attr(exc, "__traceback__", vm.ctx.new_list(new_tb))
+                .unwrap();
+        }
+    }
+    exc.clone()
+}
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -15,6 +15,7 @@ use crate::bytecode;
 use crate::frame::{ExecutionResult, Frame, FrameRef, Scope};
 use crate::frozen;
 use crate::function::PyFuncArgs;
+use crate::import;
 use crate::obj::objbool;
 use crate::obj::objbuiltinfunc::PyBuiltinFunction;
 use crate::obj::objcode::{PyCode, PyCodeRef};
@@ -305,30 +306,33 @@ impl VirtualMachine {
 
     pub fn import(&self, module: &str, from_list: &PyObjectRef, level: usize) -> PyResult {
         let sys_modules = self.get_attribute(self.sys_module.clone(), "modules")?;
-        sys_modules.get_item(module.to_string(), self).or_else(|_| {
-            let import_func = self
-                .get_attribute(self.builtins.clone(), "__import__")
-                .map_err(|_| self.new_import_error("__import__ not found".to_string()))?;
-
-            let (locals, globals) = if let Some(frame) = self.current_frame() {
-                (
-                    frame.scope.get_locals().into_object(),
-                    frame.scope.globals.clone().into_object(),
+        sys_modules
+            .get_item(module.to_string(), self)
+            .or_else(|_| {
+                let import_func = self
+                    .get_attribute(self.builtins.clone(), "__import__")
+                    .map_err(|_| self.new_import_error("__import__ not found".to_string()))?;
+
+                let (locals, globals) = if let Some(frame) = self.current_frame() {
+                    (
+                        frame.scope.get_locals().into_object(),
+                        frame.scope.globals.clone().into_object(),
+                    )
+                } else {
+                    (self.get_none(), self.get_none())
+                };
+                self.invoke(
+                    import_func,
+                    vec![
+                        self.ctx.new_str(module.to_string()),
+                        globals,
+                        locals,
+                        from_list.clone(),
+                        self.ctx.new_int(level),
+                    ],
                 )
-            } else {
-                (self.get_none(), self.get_none())
-            };
-            self.invoke(
-                import_func,
-                vec![
-                    self.ctx.new_str(module.to_string()),
-                    globals,
-                    locals,
-                    from_list.clone(),
-                    self.ctx.new_int(level),
-                ],
-            )
-        })
+            })
+            .map_err(|exc| import::remove_importlib_frames(self, &exc))
     }
 
     /// Determines if `obj` is an instance of `cls`, either directly, indirectly or virtually via
