RustPython · coolreader18 · Jul 27, 2019 · Jul 27, 2019 · Jul 27, 2019
diff --git a/bytecode/src/bytecode.rs b/bytecode/src/bytecode.rs
@@ -476,3 +476,8 @@ impl fmt::Debug for CodeObject {
         )
     }
 }
+
+pub struct FrozenModule {
+    pub code: CodeObject,
+    pub package: bool,
+}
diff --git a/derive/src/compile_bytecode.rs b/derive/src/compile_bytecode.rs
@@ -17,7 +17,7 @@ use crate::{extract_spans, Diagnostic};
 use bincode;
 use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::quote;
-use rustpython_bytecode::bytecode::CodeObject;
+use rustpython_bytecode::bytecode::{CodeObject, FrozenModule};
 use rustpython_compiler::compile;
 use std::collections::HashMap;
 use std::env;
@@ -52,7 +52,7 @@ impl CompilationSource {
         &self,
         mode: &compile::Mode,
         module_name: String,
-    ) -> Result<HashMap<String, CodeObject>, Diagnostic> {
+    ) -> Result<HashMap<String, FrozenModule>, Diagnostic> {
         Ok(match &self.kind {
             CompilationSourceKind::File(rel_path) => {
                 let mut path = PathBuf::from(
@@ -65,10 +65,20 @@ impl CompilationSource {
                         format!("Error reading file {:?}: {}", path, err),
                     )
                 })?;
-                hashmap! {module_name.clone() => self.compile_string(&source, mode, module_name.clone())?}
+                hashmap! {
+                    module_name.clone() => FrozenModule {
+                        code: self.compile_string(&source, mode, module_name.clone())?,
+                        package: false,
+                    },
+                }
             }
             CompilationSourceKind::SourceCode(code) => {
-                hashmap! {module_name.clone() => self.compile_string(code, mode, module_name.clone())?}
+                hashmap! {
+                    module_name.clone() => FrozenModule {
+                        code: self.compile_string(code, mode, module_name.clone())?,
+                        package: false,
+                    },
+                }
             }
             CompilationSourceKind::Dir(rel_path) => {
                 let mut path = PathBuf::from(
@@ -85,7 +95,7 @@ impl CompilationSource {
         path: &Path,
         parent: String,
         mode: &compile::Mode,
-    ) -> Result<HashMap<String, CodeObject>, Diagnostic> {
+    ) -> Result<HashMap<String, FrozenModule>, Diagnostic> {
         let mut code_map = HashMap::new();
         let paths = fs::read_dir(&path).map_err(|err| {
             Diagnostic::spans_error(self.span, format!("Error listing dir {:?}: {}", path, err))
@@ -95,11 +105,13 @@ impl CompilationSource {
                 Diagnostic::spans_error(self.span, format!("Failed to list file: {}", err))
             })?;
             let path = path.path();
-            let file_name = path.file_name().unwrap().to_str().unwrap();
+            let file_name = path.file_name().unwrap().to_str().ok_or_else(|| {
+                Diagnostic::spans_error(self.span, format!("Invalid UTF-8 in file name {:?}", path))
+            })?;
             if path.is_dir() {
                 code_map.extend(self.compile_dir(
                     &path,
-                    format!("{}{}.", parent, file_name),
+                    format!("{}{}", parent, file_name),
                     mode,
                 )?);
             } else if file_name.ends_with(".py") {
@@ -109,11 +121,21 @@ impl CompilationSource {
                         format!("Error reading file {:?}: {}", path, err),
                     )
                 })?;
-                let file_name_splitte: Vec<&str> = file_name.splitn(2, '.').collect();
-                let module_name = format!("{}{}", parent, file_name_splitte[0]);
+                let stem = path.file_stem().unwrap().to_str().unwrap();
+                let is_init = stem == "__init__";
+                let module_name = if is_init {
+                    parent.clone()
+                } else if parent.is_empty() {
+                    stem.to_string()
+                } else {
+                    format!("{}.{}", parent, stem)
+                };
                 code_map.insert(
                     module_name.clone(),
-                    self.compile_string(&source, mode, module_name)?,
+                    FrozenModule {
+                        code: self.compile_string(&source, mode, module_name)?,
+                        package: is_init,
+                    },
                 );
             }
         }
@@ -128,7 +150,7 @@ struct PyCompileInput {
 }
 
 impl PyCompileInput {
-    fn compile(&self) -> Result<HashMap<String, CodeObject>, Diagnostic> {
+    fn compile(&self) -> Result<HashMap<String, FrozenModule>, Diagnostic> {
         let mut module_name = None;
         let mut mode = None;
         let mut source: Option<CompilationSource> = None;
@@ -225,13 +247,21 @@ pub fn impl_py_compile_bytecode(input: TokenStream2) -> Result<TokenStream2, Dia
 
     let code_map = input.compile()?;
 
-    let modules = code_map.iter().map(|(module_name, code_obj)| {
-        let module_name = LitStr::new(&module_name, Span::call_site());
-        let bytes = bincode::serialize(&code_obj).expect("Failed to serialize");
-        let bytes = LitByteStr::new(&bytes, Span::call_site());
-        quote! {  #module_name.into() => bincode::deserialize::<::rustpython_vm::bytecode::CodeObject>(#bytes)
-                .expect("Deserializing CodeObject failed") }
-    });
+    let modules = code_map
+        .into_iter()
+        .map(|(module_name, FrozenModule { code, package })| {
+            let module_name = LitStr::new(&module_name, Span::call_site());
+            let bytes = bincode::serialize(&code).expect("Failed to serialize");
+            let bytes = LitByteStr::new(&bytes, Span::call_site());
+            quote! {
+                #module_name.into() => ::rustpython_vm::bytecode::FrozenModule {
+                    code: bincode::deserialize::<::rustpython_vm::bytecode::CodeObject>(
+                        #bytes
+                    ).expect("Deserializing CodeObject failed"),
+                    package: #package,
+                }
+            }
+        });
 
     let output = quote! {
         ({

diff --git a/vm/src/frozen.rs b/vm/src/frozen.rs
@@ -1,7 +1,7 @@
-use crate::bytecode::CodeObject;
+use crate::bytecode::FrozenModule;
 use std::collections::HashMap;
 
-pub fn get_module_inits() -> HashMap<String, CodeObject> {
+pub fn get_module_inits() -> HashMap<String, FrozenModule> {
     let mut modules = HashMap::new();
     modules.extend(py_compile_bytecode!(
         source = "initialized = True; print(\"Hello world!\")\n",

diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -43,7 +43,7 @@ pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult {
         .borrow()
         .get(module_name)
         .ok_or_else(|| vm.new_import_error(format!("Cannot import frozen module {}", module_name)))
-        .and_then(|frozen| import_codeobj(vm, module_name, frozen.clone(), false))
+        .and_then(|frozen| import_codeobj(vm, module_name, frozen.code.clone(), false))
 }
 
 pub fn import_builtin(vm: &VirtualMachine, module_name: &str) -> PyResult {

diff --git a/vm/src/stdlib/imp.rs b/vm/src/stdlib/imp.rs
@@ -57,7 +57,7 @@ fn imp_get_frozen_object(name: PyStringRef, vm: &VirtualMachine) -> PyResult<PyC
         .borrow()
         .get(name.as_str())
         .map(|frozen| {
-            let mut frozen = frozen.clone();
+            let mut frozen = frozen.code.clone();
             frozen.source_path = format!("frozen {}", name.as_str());
             PyCode::new(frozen)
         })
@@ -70,9 +70,14 @@ fn imp_init_frozen(name: PyStringRef, vm: &VirtualMachine) -> PyResult {
     import::import_frozen(vm, name.as_str())
 }
 
-fn imp_is_frozen_package(_name: PyStringRef, _vm: &VirtualMachine) -> bool {
-    // TODO: Support frozen package.
-    false
+fn imp_is_frozen_package(name: PyStringRef, vm: &VirtualMachine) -> PyResult<bool> {
+    vm.frozen
+        .borrow()
+        .get(name.as_str())
+        .map(|frozen| frozen.package)
+        .ok_or_else(|| {
+            vm.new_import_error(format!("No such frozen object named {}", name.as_str()))
+        })
 }
 
 fn imp_fix_co_filename(_code: PyObjectRef, _path: PyStringRef, _vm: &VirtualMachine) {

diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -56,7 +56,7 @@ pub struct VirtualMachine {
     pub frames: RefCell<Vec<FrameRef>>,
     pub wasm_id: Option<String>,
     pub exceptions: RefCell<Vec<PyObjectRef>>,
-    pub frozen: RefCell<HashMap<String, bytecode::CodeObject>>,
+    pub frozen: RefCell<HashMap<String, bytecode::FrozenModule>>,
     pub import_func: RefCell<PyObjectRef>,
     pub profile_func: RefCell<PyObjectRef>,
     pub trace_func: RefCell<PyObjectRef>,

diff --git a/wasm/lib/src/convert.rs b/wasm/lib/src/convert.rs
@@ -22,13 +22,9 @@ pub fn py_err_to_js_err(vm: &VirtualMachine, py_err: &PyObjectRef) -> JsValue {
             }
         };
     }
-    let msg = match vm
-        .get_attribute(py_err.clone(), "msg")
-        .ok()
-        .and_then(|msg| vm.to_pystr(&msg).ok())
-    {
-        Some(msg) => msg,
-        None => return js_sys::Error::new("error getting error").into(),
+    let msg = match vm.to_pystr(py_err) {
+        Ok(msg) => msg,
+        Err(_) => return js_sys::Error::new("error getting error").into(),
     };
     let js_err = map_exceptions!(py_err,& msg, {
         // TypeError is sort of a catch-all for "this value isn't what I thought it was like"