diff --git a/crates/compiler-core/src/frozen.rs b/crates/compiler-core/src/frozen.rs index a79569ad7f..34a626e2b1 100644 --- a/crates/compiler-core/src/frozen.rs +++ b/crates/compiler-core/src/frozen.rs @@ -1,11 +1,14 @@ use crate::bytecode::*; use crate::marshal::{self, Read, ReadBorrowed, Write}; -/// A frozen module. Holds a frozen code object and whether it is part of a package +/// A frozen module. Holds a frozen code object and whether it is part of a package. +/// The `origname` type is generic to allow either static names (runtime) or +/// borrowed/owned names during compile-time freezing. #[derive(Copy, Clone)] -pub struct FrozenModule { +pub struct FrozenModule { pub code: FrozenCodeObject, pub package: bool, + pub origname: N, } #[derive(Copy, Clone)] @@ -55,7 +58,7 @@ impl + ?Sized> FrozenLib { } impl<'a, B: AsRef<[u8]> + ?Sized> IntoIterator for &'a FrozenLib { - type Item = (&'a str, FrozenModule<&'a [u8]>); + type Item = (&'a str, FrozenModule<&'a [u8], &'a str>); type IntoIter = FrozenModulesIter<'a>; fn into_iter(self) -> Self::IntoIter { @@ -69,7 +72,7 @@ pub struct FrozenModulesIter<'a> { } impl<'a> Iterator for FrozenModulesIter<'a> { - type Item = (&'a str, FrozenModule<&'a [u8]>); + type Item = (&'a str, FrozenModule<&'a [u8], &'a str>); fn next(&mut self) -> Option { if self.remaining > 0 { @@ -90,21 +93,30 @@ impl ExactSizeIterator for FrozenModulesIter<'_> {} fn read_entry<'a>( rdr: &mut &'a [u8], -) -> Result<(&'a str, FrozenModule<&'a [u8]>), marshal::MarshalError> { +) -> Result<(&'a str, FrozenModule<&'a [u8], &'a str>), marshal::MarshalError> { let len = rdr.read_u32()?; let name = rdr.read_str_borrow(len)?; let len = rdr.read_u32()?; let code_slice = rdr.read_slice_borrow(len)?; let code = FrozenCodeObject { bytes: code_slice }; let package = rdr.read_u8()? != 0; - Ok((name, FrozenModule { code, package })) + let len = rdr.read_u32()?; + let origname = rdr.read_str_borrow(len)?; + Ok(( + name, + FrozenModule { + code, + package, + origname, + }, + )) } impl FrozenLib> { /// Encode the given iterator of frozen modules into a compressed vector of bytes - pub fn encode<'a, I, B: AsRef<[u8]>>(lib: I) -> Self + pub fn encode<'a, I, B: AsRef<[u8]>, N: AsRef>(lib: I) -> Self where - I: IntoIterator), IntoIter: ExactSizeIterator + Clone>, + I: IntoIterator), IntoIter: ExactSizeIterator + Clone>, { let iter = lib.into_iter(); let mut bytes = Vec::new(); @@ -113,9 +125,9 @@ impl FrozenLib> { } } -fn write_lib<'a, B: AsRef<[u8]>>( +fn write_lib<'a, B: AsRef<[u8]>, N: AsRef>( buf: &mut Vec, - lib: impl ExactSizeIterator)>, + lib: impl ExactSizeIterator)>, ) { marshal::write_len(buf, lib.len()); for (name, module) in lib { @@ -123,8 +135,13 @@ fn write_lib<'a, B: AsRef<[u8]>>( } } -fn write_entry(buf: &mut Vec, name: &str, module: FrozenModule>) { +fn write_entry>( + buf: &mut Vec, + name: &str, + module: FrozenModule, N>, +) { marshal::write_vec(buf, name.as_bytes()); marshal::write_vec(buf, module.code.bytes.as_ref()); buf.write_u8(module.package as u8); + marshal::write_vec(buf, module.origname.as_ref().as_bytes()); } diff --git a/crates/derive-impl/src/compile_bytecode.rs b/crates/derive-impl/src/compile_bytecode.rs index cdcc89b998..80fafa4754 100644 --- a/crates/derive-impl/src/compile_bytecode.rs +++ b/crates/derive-impl/src/compile_bytecode.rs @@ -45,6 +45,7 @@ enum CompilationSourceKind { struct CompiledModule { code: CodeObject, package: bool, + origname: String, } struct CompilationSource { @@ -91,12 +92,17 @@ impl CompilationSource { mode, compiler, ), - _ => Ok(hashmap! { - module_name.clone() => CompiledModule { - code: self.compile_single(mode, module_name, compiler)?, - package: false, - }, - }), + _ => { + let origname = module_name.clone(); + let code = self.compile_single(mode, module_name, compiler)?; + Ok(hashmap! { + origname.clone() => CompiledModule { + code, + package: false, + origname, + }, + }) + } } } @@ -220,11 +226,13 @@ impl CompilationSource { Err(e) => return Err(e), }; + let origname = module_name.clone(); code_map.insert( module_name, CompiledModule { code, package: is_init, + origname, }, ); } @@ -368,6 +376,7 @@ pub fn impl_py_freeze( let v = frozen::FrozenModule { code: frozen::FrozenCodeObject::encode(&v.code), package: v.package, + origname: &*v.origname, }; (&**k, v) })); diff --git a/crates/vm/src/import.rs b/crates/vm/src/import.rs index 39748655e0..0b894863f5 100644 --- a/crates/vm/src/import.rs +++ b/crates/vm/src/import.rs @@ -80,11 +80,15 @@ 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)?; + module.set_attr("__origname__", vm.ctx.new_str(frozen.origname), vm)?; Ok(module) } diff --git a/crates/vm/src/stdlib/imp.rs b/crates/vm/src/stdlib/imp.rs index 76b3bfd124..81b711f851 100644 --- a/crates/vm/src/stdlib/imp.rs +++ b/crates/vm/src/stdlib/imp.rs @@ -188,7 +188,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(info.origname); Ok(Some((None, info.package, origname))) } diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index 3409245405..b245104f0d 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -983,6 +983,10 @@ impl AsRef for VirtualMachine { } } +use std::sync::OnceLock; + +static FROZEN_ORIGNAME_ALIASES: OnceLock> = OnceLock::new(); + fn core_frozen_inits() -> impl Iterator { let iter = std::iter::empty(); macro_rules! ext_modules { @@ -1018,7 +1022,22 @@ fn core_frozen_inits() -> impl Iterator { crate_name = "rustpython_compiler_core" ); - iter + let aliases = FROZEN_ORIGNAME_ALIASES.get_or_init(|| { + HashMap::from([ + ("_frozen_importlib", "importlib._bootstrap"), + ( + "_frozen_importlib_external", + "importlib._bootstrap_external", + ), + ]) + }); + + iter.map(|(name, mut module)| { + if let Some(origname) = aliases.get(name) { + module.origname = *origname; + } + (name, module) + }) } #[test] @@ -1046,3 +1065,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", + ); + }); +}