Skip to content

Commit fc97bd9

Browse files
coolreader18coolreader18
authored
Merge pull request #1185 from RustPython/coolreader18/frozen-packages
Add frozen package support
2 parents a02e131 + ec6c016 commit fc97bd9

File tree

7 files changed

+69
-33
lines changed

7 files changed

+69
-33
lines changed

‎bytecode/src/bytecode.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -476,3 +476,8 @@ impl fmt::Debug for CodeObject {
476476
)
477477
}
478478
}
479+
480+
pub struct FrozenModule {
481+
pub code: CodeObject,
482+
pub package: bool,
483+
}

‎derive/src/compile_bytecode.rs‎

Lines changed: 48 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ use crate::{extract_spans, Diagnostic};
1717
use bincode;
1818
use proc_macro2::{Span, TokenStream as TokenStream2};
1919
use quote::quote;
20-
use rustpython_bytecode::bytecode::CodeObject;
20+
use rustpython_bytecode::bytecode::{CodeObject, FrozenModule};
2121
use rustpython_compiler::compile;
2222
use std::collections::HashMap;
2323
use std::env;
@@ -52,7 +52,7 @@ impl CompilationSource {
5252
&self,
5353
mode: &compile::Mode,
5454
module_name: String,
55-
) -> Result<HashMap<String, CodeObject>, Diagnostic> {
55+
) -> Result<HashMap<String, FrozenModule>, Diagnostic> {
5656
Ok(match &self.kind {
5757
CompilationSourceKind::File(rel_path) => {
5858
let mut path = PathBuf::from(
@@ -65,10 +65,20 @@ impl CompilationSource {
6565
format!("Error reading file {:?}: {}", path, err),
6666
)
6767
})?;
68-
hashmap! {module_name.clone() => self.compile_string(&source, mode, module_name.clone())?}
68+
hashmap! {
69+
module_name.clone() => FrozenModule {
70+
code: self.compile_string(&source, mode, module_name.clone())?,
71+
package: false,
72+
},
73+
}
6974
}
7075
CompilationSourceKind::SourceCode(code) => {
71-
hashmap! {module_name.clone() => self.compile_string(code, mode, module_name.clone())?}
76+
hashmap! {
77+
module_name.clone() => FrozenModule {
78+
code: self.compile_string(code, mode, module_name.clone())?,
79+
package: false,
80+
},
81+
}
7282
}
7383
CompilationSourceKind::Dir(rel_path) => {
7484
let mut path = PathBuf::from(
@@ -85,7 +95,7 @@ impl CompilationSource {
8595
path: &Path,
8696
parent: String,
8797
mode: &compile::Mode,
88-
) -> Result<HashMap<String, CodeObject>, Diagnostic> {
98+
) -> Result<HashMap<String, FrozenModule>, Diagnostic> {
8999
let mut code_map = HashMap::new();
90100
let paths = fs::read_dir(&path).map_err(|err| {
91101
Diagnostic::spans_error(self.span, format!("Error listing dir {:?}: {}", path, err))
@@ -95,11 +105,13 @@ impl CompilationSource {
95105
Diagnostic::spans_error(self.span, format!("Failed to list file: {}", err))
96106
})?;
97107
let path = path.path();
98-
let file_name = path.file_name().unwrap().to_str().unwrap();
108+
let file_name = path.file_name().unwrap().to_str().ok_or_else(|| {
109+
Diagnostic::spans_error(self.span, format!("Invalid UTF-8 in file name {:?}", path))
110+
})?;
99111
if path.is_dir() {
100112
code_map.extend(self.compile_dir(
101113
&path,
102-
format!("{}{}.", parent, file_name),
114+
format!("{}{}", parent, file_name),
103115
mode,
104116
)?);
105117
} else if file_name.ends_with(".py") {
@@ -109,11 +121,21 @@ impl CompilationSource {
109121
format!("Error reading file {:?}: {}", path, err),
110122
)
111123
})?;
112-
let file_name_splitte: Vec<&str> = file_name.splitn(2, '.').collect();
113-
let module_name = format!("{}{}", parent, file_name_splitte[0]);
124+
let stem = path.file_stem().unwrap().to_str().unwrap();
125+
let is_init = stem == "__init__";
126+
let module_name = if is_init {
127+
parent.clone()
128+
} else if parent.is_empty() {
129+
stem.to_string()
130+
} else {
131+
format!("{}.{}", parent, stem)
132+
};
114133
code_map.insert(
115134
module_name.clone(),
116-
self.compile_string(&source, mode, module_name)?,
135+
FrozenModule {
136+
code: self.compile_string(&source, mode, module_name)?,
137+
package: is_init,
138+
},
117139
);
118140
}
119141
}
@@ -128,7 +150,7 @@ struct PyCompileInput {
128150
}
129151

130152
impl PyCompileInput {
131-
fn compile(&self) -> Result<HashMap<String, CodeObject>, Diagnostic> {
153+
fn compile(&self) -> Result<HashMap<String, FrozenModule>, Diagnostic> {
132154
let mut module_name = None;
133155
let mut mode = None;
134156
let mut source: Option<CompilationSource> = None;
@@ -225,13 +247,21 @@ pub fn impl_py_compile_bytecode(input: TokenStream2) -> Result<TokenStream2, Dia
225247

226248
let code_map = input.compile()?;
227249

228-
let modules = code_map.iter().map(|(module_name, code_obj)| {
229-
let module_name = LitStr::new(&module_name, Span::call_site());
230-
let bytes = bincode::serialize(&code_obj).expect("Failed to serialize");
231-
let bytes = LitByteStr::new(&bytes, Span::call_site());
232-
quote! { #module_name.into() => bincode::deserialize::<::rustpython_vm::bytecode::CodeObject>(#bytes)
233-
.expect("Deserializing CodeObject failed") }
234-
});
250+
let modules = code_map
251+
.into_iter()
252+
.map(|(module_name, FrozenModule { code, package })| {
253+
let module_name = LitStr::new(&module_name, Span::call_site());
254+
let bytes = bincode::serialize(&code).expect("Failed to serialize");
255+
let bytes = LitByteStr::new(&bytes, Span::call_site());
256+
quote! {
257+
#module_name.into() => ::rustpython_vm::bytecode::FrozenModule {
258+
code: bincode::deserialize::<::rustpython_vm::bytecode::CodeObject>(
259+
#bytes
260+
).expect("Deserializing CodeObject failed"),
261+
package: #package,
262+
}
263+
}
264+
});
235265

236266
let output = quote! {
237267
({

‎vm/src/frozen.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
use crate::bytecode::CodeObject;
1+
use crate::bytecode::FrozenModule;
22
use std::collections::HashMap;
33

4-
pub fn get_module_inits() -> HashMap<String, CodeObject> {
4+
pub fn get_module_inits() -> HashMap<String, FrozenModule> {
55
let mut modules = HashMap::new();
66
modules.extend(py_compile_bytecode!(
77
source = "initialized = True; print(\"Hello world!\")\n",

‎vm/src/import.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult {
4343
.borrow()
4444
.get(module_name)
4545
.ok_or_else(|| vm.new_import_error(format!("Cannot import frozen module {}", module_name)))
46-
.and_then(|frozen| import_codeobj(vm, module_name, frozen.clone(), false))
46+
.and_then(|frozen| import_codeobj(vm, module_name, frozen.code.clone(), false))
4747
}
4848

4949
pub fn import_builtin(vm: &VirtualMachine, module_name: &str) -> PyResult {

‎vm/src/stdlib/imp.rs‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ fn imp_get_frozen_object(name: PyStringRef, vm: &VirtualMachine) -> PyResult<PyC
5757
.borrow()
5858
.get(name.as_str())
5959
.map(|frozen| {
60-
let mut frozen = frozen.clone();
60+
let mut frozen = frozen.code.clone();
6161
frozen.source_path = format!("frozen {}", name.as_str());
6262
PyCode::new(frozen)
6363
})
@@ -70,9 +70,14 @@ fn imp_init_frozen(name: PyStringRef, vm: &VirtualMachine) -> PyResult {
7070
import::import_frozen(vm, name.as_str())
7171
}
7272

73-
fn imp_is_frozen_package(_name: PyStringRef, _vm: &VirtualMachine) -> bool {
74-
// TODO: Support frozen package.
75-
false
73+
fn imp_is_frozen_package(name: PyStringRef, vm: &VirtualMachine) -> PyResult<bool> {
74+
vm.frozen
75+
.borrow()
76+
.get(name.as_str())
77+
.map(|frozen| frozen.package)
78+
.ok_or_else(|| {
79+
vm.new_import_error(format!("No such frozen object named {}", name.as_str()))
80+
})
7681
}
7782

7883
fn imp_fix_co_filename(_code: PyObjectRef, _path: PyStringRef, _vm: &VirtualMachine) {

‎vm/src/vm.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ pub struct VirtualMachine {
5656
pub frames: RefCell<Vec<FrameRef>>,
5757
pub wasm_id: Option<String>,
5858
pub exceptions: RefCell<Vec<PyObjectRef>>,
59-
pub frozen: RefCell<HashMap<String, bytecode::CodeObject>>,
59+
pub frozen: RefCell<HashMap<String, bytecode::FrozenModule>>,
6060
pub import_func: RefCell<PyObjectRef>,
6161
pub profile_func: RefCell<PyObjectRef>,
6262
pub trace_func: RefCell<PyObjectRef>,

‎wasm/lib/src/convert.rs‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,13 +22,9 @@ pub fn py_err_to_js_err(vm: &VirtualMachine, py_err: &PyObjectRef) -> JsValue {
2222
}
2323
};
2424
}
25-
let msg = match vm
26-
.get_attribute(py_err.clone(), "msg")
27-
.ok()
28-
.and_then(|msg| vm.to_pystr(&msg).ok())
29-
{
30-
Some(msg) => msg,
31-
None => return js_sys::Error::new("error getting error").into(),
25+
let msg = match vm.to_pystr(py_err) {
26+
Ok(msg) => msg,
27+
Err(_) => return js_sys::Error::new("error getting error").into(),
3228
};
3329
let js_err = map_exceptions!(py_err,& msg, {
3430
// TypeError is sort of a catch-all for "this value isn't what I thought it was like"

0 commit comments

Comments
 (0)