diff --git a/vm/src/import.rs b/vm/src/import.rs index 28727b6bfd..2c2ac87289 100644 --- 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 index bd0beba5ba..f117b83f31 100644 --- 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