diff --git a/crates/vm/src/import.rs b/crates/vm/src/import.rs index f970a335aa5..d67bf59b35f 100644 --- a/crates/vm/src/import.rs +++ b/crates/vm/src/import.rs @@ -5,7 +5,7 @@ use crate::{ builtins::{PyCode, list, traceback::PyTraceback}, exceptions::types::PyBaseException, scope::Scope, - vm::{VirtualMachine, thread}, + vm::{VirtualMachine, resolve_frozen_alias, thread}, }; pub(crate) fn init_importlib_base(vm: &mut VirtualMachine) -> PyResult { @@ -69,11 +69,16 @@ pub fn make_frozen(vm: &VirtualMachine, name: &str) -> PyResult> { } pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult { - let frozen = make_frozen(vm, module_name)?; - let module = import_code_obj(vm, module_name, frozen, false)?; + let frozen = vm.state.frozen.get(module_name).ok_or_else(|| { + vm.new_import_error( + format!("No such frozen object named {module_name}"), + vm.ctx.new_str(module_name), + ) + })?; + let module = import_code_obj(vm, module_name, vm.ctx.new_code(frozen.code), false)?; debug_assert!(module.get_attr(identifier!(vm, __name__), vm).is_ok()); - // TODO: give a correct origname here - module.set_attr("__origname__", vm.ctx.new_str(module_name.to_owned()), vm)?; + let origname = resolve_frozen_alias(module_name); + module.set_attr("__origname__", vm.ctx.new_str(origname), vm)?; Ok(module) } diff --git a/crates/vm/src/stdlib/imp.rs b/crates/vm/src/stdlib/imp.rs index df32b0c0068..ad2cede975c 100644 --- a/crates/vm/src/stdlib/imp.rs +++ b/crates/vm/src/stdlib/imp.rs @@ -2,6 +2,8 @@ use crate::frozen::FrozenModule; use crate::{VirtualMachine, builtins::PyBaseExceptionRef}; pub(crate) use _imp::make_module; +pub use crate::vm::resolve_frozen_alias; + #[cfg(feature = "threading")] #[pymodule(sub)] mod lock { @@ -191,7 +193,7 @@ mod _imp { Err(e) => return Err(e.to_pyexception(name.as_str(), vm)), }; - let origname = name; // FIXME: origname != name + let origname = vm.ctx.new_str(super::resolve_frozen_alias(name.as_str())); Ok(Some((None, info.package, origname))) } diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index b156e30b738..2eef479d1b4 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -1001,6 +1001,16 @@ impl AsRef for VirtualMachine { } } +/// Resolve frozen module alias to its original name. +/// Returns the original module name if an alias exists, otherwise returns the input name. +pub fn resolve_frozen_alias(name: &str) -> &str { + match name { + "_frozen_importlib" => "importlib._bootstrap", + "_frozen_importlib_external" => "importlib._bootstrap_external", + _ => name, + } +} + fn core_frozen_inits() -> impl Iterator { let iter = core::iter::empty(); macro_rules! ext_modules { @@ -1064,3 +1074,26 @@ fn test_nested_frozen() { } }) } + +#[test] +fn frozen_origname_matches() { + use rustpython_vm as vm; + + vm::Interpreter::with_init(Default::default(), |_vm| {}).enter(|vm| { + let check = |name, expected| { + let module = import::import_frozen(vm, name).unwrap(); + let origname: PyStrRef = module + .get_attr("__origname__", vm) + .unwrap() + .try_into_value(vm) + .unwrap(); + assert_eq!(origname.as_str(), expected); + }; + + check("_frozen_importlib", "importlib._bootstrap"); + check( + "_frozen_importlib_external", + "importlib._bootstrap_external", + ); + }); +}