Skip to content

Commit 72f314d

Browse files
youknowoneyouknowone
committed
reimpl origname
1 parent 873c7f2 commit 72f314d

File tree

5 files changed

+33
-65
lines changed

5 files changed

+33
-65
lines changed

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

Lines changed: 11 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,11 @@
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.
5-
/// The `origname` type is generic to allow either static names (runtime) or
6-
/// borrowed/owned names during compile-time freezing.
4+
/// A frozen module. Holds a frozen code object and whether it is part of a package
75
#[derive(Copy, Clone)]
8-
pub struct FrozenModule<B = &'static [u8], N = &'static str> {
6+
pub struct FrozenModule<B = &'static [u8]> {
97
pub code: FrozenCodeObject<B>,
108
pub package: bool,
11-
pub origname: N,
129
}
1310

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

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

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

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

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

9491
fn read_entry<'a>(
9592
rdr: &mut &'a [u8],
96-
) -> Result<(&'a str, FrozenModule<&'a [u8], &'a str>), marshal::MarshalError> {
93+
) -> Result<(&'a str, FrozenModule<&'a [u8]>), marshal::MarshalError> {
9794
let len = rdr.read_u32()?;
9895
let name = rdr.read_str_borrow(len)?;
9996
let len = rdr.read_u32()?;
10097
let code_slice = rdr.read_slice_borrow(len)?;
10198
let code = FrozenCodeObject { bytes: code_slice };
10299
let package = rdr.read_u8()? != 0;
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-
))
100+
Ok((name, FrozenModule { code, package }))
113101
}
114102

115103
impl FrozenLib<Vec<u8>> {
116104
/// Encode the given iterator of frozen modules into a compressed vector of bytes
117-
pub fn encode<'a, I, B: AsRef<[u8]>, N: AsRef<str>>(lib: I) -> Self
105+
pub fn encode<'a, I, B: AsRef<[u8]>>(lib: I) -> Self
118106
where
119-
I: IntoIterator<Item = (&'a str, FrozenModule<B, N>), IntoIter: ExactSizeIterator + Clone>,
107+
I: IntoIterator<Item = (&'a str, FrozenModule<B>), IntoIter: ExactSizeIterator + Clone>,
120108
{
121109
let iter = lib.into_iter();
122110
let mut bytes = Vec::new();
@@ -125,23 +113,18 @@ impl FrozenLib<Vec<u8>> {
125113
}
126114
}
127115

128-
fn write_lib<'a, B: AsRef<[u8]>, N: AsRef<str>>(
116+
fn write_lib<'a, B: AsRef<[u8]>>(
129117
buf: &mut Vec<u8>,
130-
lib: impl ExactSizeIterator<Item = (&'a str, FrozenModule<B, N>)>,
118+
lib: impl ExactSizeIterator<Item = (&'a str, FrozenModule<B>)>,
131119
) {
132120
marshal::write_len(buf, lib.len());
133121
for (name, module) in lib {
134122
write_entry(buf, name, module);
135123
}
136124
}
137125

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

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

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

5150
struct CompilationSource {
@@ -92,17 +91,12 @@ impl CompilationSource {
9291
mode,
9392
compiler,
9493
),
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-
}
94+
_ => Ok(hashmap! {
95+
module_name.clone() => CompiledModule {
96+
code: self.compile_single(mode, module_name, compiler)?,
97+
package: false,
98+
},
99+
}),
106100
}
107101
}
108102

@@ -226,13 +220,11 @@ impl CompilationSource {
226220
Err(e) => return Err(e),
227221
};
228222

229-
let origname = module_name.clone();
230223
code_map.insert(
231224
module_name,
232225
CompiledModule {
233226
code,
234227
package: is_init,
235-
origname,
236228
},
237229
);
238230
}
@@ -376,7 +368,6 @@ pub fn impl_py_freeze(
376368
let v = frozen::FrozenModule {
377369
code: frozen::FrozenCodeObject::encode(&v.code),
378370
package: v.package,
379-
origname: &*v.origname,
380371
};
381372
(&**k, v)
382373
}));

‎crates/vm/src/import.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use crate::{
55
builtins::{PyCode, list, traceback::PyTraceback},
66
exceptions::types::PyBaseException,
77
scope::Scope,
8-
vm::{VirtualMachine, thread},
8+
vm::{VirtualMachine, resolve_frozen_alias, thread},
99
};
1010

1111
pub(crate) fn init_importlib_base(vm: &mut VirtualMachine) -> PyResult<PyObjectRef> {
@@ -77,7 +77,8 @@ pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult {
7777
})?;
7878
let module = import_code_obj(vm, module_name, vm.ctx.new_code(frozen.code), false)?;
7979
debug_assert!(module.get_attr(identifier!(vm, __name__), vm).is_ok());
80-
module.set_attr("__origname__", vm.ctx.new_str(frozen.origname), vm)?;
80+
let origname = resolve_frozen_alias(module_name);
81+
module.set_attr("__origname__", vm.ctx.new_str(origname), vm)?;
8182
Ok(module)
8283
}
8384

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ use crate::frozen::FrozenModule;
22
use crate::{VirtualMachine, builtins::PyBaseExceptionRef};
33
pub(crate) use _imp::make_module;
44

5+
pub use crate::vm::resolve_frozen_alias;
6+
57
#[cfg(feature = "threading")]
68
#[pymodule(sub)]
79
mod lock {
@@ -191,7 +193,7 @@ mod _imp {
191193
Err(e) => return Err(e.to_pyexception(name.as_str(), vm)),
192194
};
193195

194-
let origname = vm.ctx.new_str(info.origname);
196+
let origname = vm.ctx.new_str(super::resolve_frozen_alias(name.as_str()));
195197
Ok(Some((None, info.package, origname)))
196198
}
197199

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

Lines changed: 10 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1001,9 +1001,15 @@ 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();
1004+
/// Resolve frozen module alias to its original name.
1005+
/// Returns the original module name if an alias exists, otherwise returns the input name.
1006+
pub fn resolve_frozen_alias(name: &str) -> &str {
1007+
match name {
1008+
"_frozen_importlib" => "importlib._bootstrap",
1009+
"_frozen_importlib_external" => "importlib._bootstrap_external",
1010+
_ => name,
1011+
}
1012+
}
10071013

10081014
fn core_frozen_inits() -> impl Iterator<Item = (&'static str, FrozenModule)> {
10091015
let iter = core::iter::empty();
@@ -1040,22 +1046,7 @@ fn core_frozen_inits() -> impl Iterator<Item = (&'static str, FrozenModule)> {
10401046
crate_name = "rustpython_compiler_core"
10411047
);
10421048

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-
iter.map(|(name, mut module)| {
1054-
if let Some(origname) = aliases.get(name) {
1055-
module.origname = *origname;
1056-
}
1057-
(name, module)
1058-
})
1049+
iter
10591050
}
10601051

10611052
#[test]

0 commit comments

Comments
 (0)