Skip to content

Commit 6af48b2

Browse files
Copilotyouknowone
Copilot
and
youknowone
committed
Fix frozen module origname metadata
Co-authored-by: youknowone <69878+youknowone@users.noreply.github.com>
1 parent 7a37e62 commit 6af48b2

5 files changed

Lines changed: 96 additions & 22 deletions

File tree

‎crates/compiler-core/src/frozen.rs‎

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
use crate::bytecode::*;
22
use crate::marshal::{self, Read, ReadBorrowed, Write};
33

4-
/// A frozen module. Holds a frozen code object and whether it is part of a package
4+
/// A frozen module. Holds a frozen code object and whether it is part of a package.
5+
/// The `origname` type is generic to allow either static names (runtime) or
6+
/// borrowed/owned names during compile-time freezing.
57
#[derive(Copy, Clone)]
6-
pub struct FrozenModule<B = &'static [u8]> {
8+
pub struct FrozenModule<B = &'static [u8], N = &'static str> {
79
pub code: FrozenCodeObject<B>,
810
pub package: bool,
11+
pub origname: N,
912
}
1013

1114
#[derive(Copy, Clone)]
@@ -55,7 +58,7 @@ impl<B: AsRef<[u8]> + ?Sized> FrozenLib<B> {
5558
}
5659

5760
impl<'a, B: AsRef<[u8]> + ?Sized> IntoIterator for &'a FrozenLib<B> {
58-
type Item = (&'a str, FrozenModule<&'a [u8]>);
61+
type Item = (&'a str, FrozenModule<&'a [u8], &'a str>);
5962
type IntoIter = FrozenModulesIter<'a>;
6063

6164
fn into_iter(self) -> Self::IntoIter {
@@ -69,7 +72,7 @@ pub struct FrozenModulesIter<'a> {
6972
}
7073

7174
impl<'a> Iterator for FrozenModulesIter<'a> {
72-
type Item = (&'a str, FrozenModule<&'a [u8]>);
75+
type Item = (&'a str, FrozenModule<&'a [u8], &'a str>);
7376

7477
fn next(&mut self) -> Option<Self::Item> {
7578
if self.remaining > 0 {
@@ -90,21 +93,30 @@ impl ExactSizeIterator for FrozenModulesIter<'_> {}
9093

9194
fn read_entry<'a>(
9295
rdr: &mut &'a [u8],
93-
) -> Result<(&'a str, FrozenModule<&'a [u8]>), marshal::MarshalError> {
96+
) -> Result<(&'a str, FrozenModule<&'a [u8], &'a str>), marshal::MarshalError> {
9497
let len = rdr.read_u32()?;
9598
let name = rdr.read_str_borrow(len)?;
9699
let len = rdr.read_u32()?;
97100
let code_slice = rdr.read_slice_borrow(len)?;
98101
let code = FrozenCodeObject { bytes: code_slice };
99102
let package = rdr.read_u8()? != 0;
100-
Ok((name, FrozenModule { code, package }))
103+
let len = rdr.read_u32()?;
104+
let origname = rdr.read_str_borrow(len)?;
105+
Ok((
106+
name,
107+
FrozenModule {
108+
code,
109+
package,
110+
origname,
111+
},
112+
))
101113
}
102114

103115
impl FrozenLib<Vec<u8>> {
104116
/// Encode the given iterator of frozen modules into a compressed vector of bytes
105-
pub fn encode<'a, I, B: AsRef<[u8]>>(lib: I) -> Self
117+
pub fn encode<'a, I, B: AsRef<[u8]>, N: AsRef<str>>(lib: I) -> Self
106118
where
107-
I: IntoIterator<Item = (&'a str, FrozenModule<B>), IntoIter: ExactSizeIterator + Clone>,
119+
I: IntoIterator<Item = (&'a str, FrozenModule<B, N>), IntoIter: ExactSizeIterator + Clone>,
108120
{
109121
let iter = lib.into_iter();
110122
let mut bytes = Vec::new();
@@ -113,18 +125,23 @@ impl FrozenLib<Vec<u8>> {
113125
}
114126
}
115127

116-
fn write_lib<'a, B: AsRef<[u8]>>(
128+
fn write_lib<'a, B: AsRef<[u8]>, N: AsRef<str>>(
117129
buf: &mut Vec<u8>,
118-
lib: impl ExactSizeIterator<Item = (&'a str, FrozenModule<B>)>,
130+
lib: impl ExactSizeIterator<Item = (&'a str, FrozenModule<B, N>)>,
119131
) {
120132
marshal::write_len(buf, lib.len());
121133
for (name, module) in lib {
122134
write_entry(buf, name, module);
123135
}
124136
}
125137

126-
fn write_entry(buf: &mut Vec<u8>, name: &str, module: FrozenModule<impl AsRef<[u8]>>) {
138+
fn write_entry<N: AsRef<str>>(
139+
buf: &mut Vec<u8>,
140+
name: &str,
141+
module: FrozenModule<impl AsRef<[u8]>, N>,
142+
) {
127143
marshal::write_vec(buf, name.as_bytes());
128144
marshal::write_vec(buf, module.code.bytes.as_ref());
129145
buf.write_u8(module.package as u8);
146+
marshal::write_vec(buf, module.origname.as_ref().as_bytes());
130147
}

‎crates/derive-impl/src/compile_bytecode.rs‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ enum CompilationSourceKind {
4545
struct CompiledModule {
4646
code: CodeObject,
4747
package: bool,
48+
origname: String,
4849
}
4950

5051
struct CompilationSource {
@@ -91,12 +92,17 @@ impl CompilationSource {
9192
mode,
9293
compiler,
9394
),
94-
_ => Ok(hashmap! {
95-
module_name.clone() => CompiledModule {
96-
code: self.compile_single(mode, module_name, compiler)?,
97-
package: false,
98-
},
99-
}),
95+
_ => {
96+
let origname = module_name.clone();
97+
let code = self.compile_single(mode, module_name, compiler)?;
98+
Ok(hashmap! {
99+
origname.clone() => CompiledModule {
100+
code,
101+
package: false,
102+
origname,
103+
},
104+
})
105+
}
100106
}
101107
}
102108

@@ -220,11 +226,13 @@ impl CompilationSource {
220226
Err(e) => return Err(e),
221227
};
222228

229+
let origname = module_name.clone();
223230
code_map.insert(
224231
module_name,
225232
CompiledModule {
226233
code,
227234
package: is_init,
235+
origname,
228236
},
229237
);
230238
}
@@ -368,6 +376,7 @@ pub fn impl_py_freeze(
368376
let v = frozen::FrozenModule {
369377
code: frozen::FrozenCodeObject::encode(&v.code),
370378
package: v.package,
379+
origname: &*v.origname,
371380
};
372381
(&**k, v)
373382
}));

‎crates/vm/src/import.rs‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -69,11 +69,15 @@ pub fn make_frozen(vm: &VirtualMachine, name: &str) -> PyResult<PyRef<PyCode>> {
6969
}
7070

7171
pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult {
72-
let frozen = make_frozen(vm, module_name)?;
73-
let module = import_code_obj(vm, module_name, frozen, false)?;
72+
let frozen = vm.state.frozen.get(module_name).ok_or_else(|| {
73+
vm.new_import_error(
74+
format!("No such frozen object named {module_name}"),
75+
vm.ctx.new_str(module_name),
76+
)
77+
})?;
78+
let module = import_code_obj(vm, module_name, vm.ctx.new_code(frozen.code), false)?;
7479
debug_assert!(module.get_attr(identifier!(vm, __name__), vm).is_ok());
75-
// TODO: give a correct origname here
76-
module.set_attr("__origname__", vm.ctx.new_str(module_name.to_owned()), vm)?;
80+
module.set_attr("__origname__", vm.ctx.new_str(frozen.origname), vm)?;
7781
Ok(module)
7882
}
7983

‎crates/vm/src/stdlib/imp.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ mod _imp {
191191
Err(e) => return Err(e.to_pyexception(name.as_str(), vm)),
192192
};
193193

194-
let origname = name; // FIXME: origname != name
194+
let origname = vm.ctx.new_str(info.origname);
195195
Ok(Some((None, info.package, origname)))
196196
}
197197

‎crates/vm/src/vm/mod.rs‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1001,6 +1001,10 @@ impl AsRef<Context> for VirtualMachine {
10011001
}
10021002
}
10031003

1004+
use std::sync::OnceLock;
1005+
1006+
static FROZEN_ORIGNAME_ALIASES: OnceLock<HashMap<&'static str, &'static str>> = OnceLock::new();
1007+
10041008
fn core_frozen_inits() -> impl Iterator<Item = (&'static str, FrozenModule)> {
10051009
let iter = core::iter::empty();
10061010
macro_rules! ext_modules {
@@ -1036,6 +1040,23 @@ fn core_frozen_inits() -> impl Iterator<Item = (&'static str, FrozenModule)> {
10361040
crate_name = "rustpython_compiler_core"
10371041
);
10381042

1043+
let aliases = FROZEN_ORIGNAME_ALIASES.get_or_init(|| {
1044+
HashMap::from([
1045+
("_frozen_importlib", "importlib._bootstrap"),
1046+
(
1047+
"_frozen_importlib_external",
1048+
"importlib._bootstrap_external",
1049+
),
1050+
])
1051+
});
1052+
1053+
let iter = iter.map(|(name, mut module)| {
1054+
if let Some(origname) = aliases.get(name) {
1055+
module.origname = *origname;
1056+
}
1057+
(name, module)
1058+
});
1059+
10391060
iter
10401061
}
10411062

@@ -1064,3 +1085,26 @@ fn test_nested_frozen() {
10641085
}
10651086
})
10661087
}
1088+
1089+
#[test]
1090+
fn frozen_origname_matches() {
1091+
use rustpython_vm as vm;
1092+
1093+
vm::Interpreter::with_init(Default::default(), |_vm| {}).enter(|vm| {
1094+
let check = |name, expected| {
1095+
let module = import::import_frozen(vm, name).unwrap();
1096+
let origname: PyStrRef = module
1097+
.get_attr("__origname__", vm)
1098+
.unwrap()
1099+
.try_into_value(vm)
1100+
.unwrap();
1101+
assert_eq!(origname.as_str(), expected);
1102+
};
1103+
1104+
check("_frozen_importlib", "importlib._bootstrap");
1105+
check(
1106+
"_frozen_importlib_external",
1107+
"importlib._bootstrap_external",
1108+
);
1109+
});
1110+
}

0 commit comments

Comments
 (0)