Merge pull request RustPython#1104 from palaviv/remove-importlib-frames

windelbouwman · web-flow · commit a099ba5f8e55 · 2019-07-06T10:58:20.000+02:00
Remove importlib frames from traceback
diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -4,8 +4,8 @@
 
 use crate::bytecode::CodeObject;
 use crate::frame::Scope;
-use crate::obj::objcode;
-use crate::pyobject::{ItemProtocol, PyResult, PyValue};
+use crate::obj::{objcode, objsequence, objstr, objtype};
+use crate::pyobject::{ItemProtocol, PyObjectRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
 #[cfg(feature = "rustpython-compiler")]
 use rustpython_compiler::compile;
@@ -82,3 +82,41 @@ pub fn import_codeobj(
     )?;
     Ok(module)
 }
+
+// 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;
+                        }
+                        if always_trim || in_importlib {
+                            false
+                        } else {
+                            true
+                        }
+                    } else {
+                        in_importlib = false;
+                        true
+                    }
+                })
+                .map(|x| x.clone())
+                .collect();
+            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
