Add set_file_attr to import_codeobj

Author: coolreader18

derive/src/compile_bytecode.rs

@@ -8,8 +8,8 @@
 //!
 //!     // the mode to compile the code in
 //!     mode = "exec", // or "eval" or "single"
-//!     // the path put into the CodeObject, defaults to `None`
-//!     source_path = "frozen",
+//!     // the path put into the CodeObject, defaults to "frozen"
+//!     module_name = "frozen",
 //! )
 //! ```
 
@@ -35,9 +35,9 @@ struct CompilationSource {
 }
 
 impl CompilationSource {
-    fn compile(self, mode: &compile::Mode, source_path: String) -> Result<CodeObject, Diagnostic> {
+    fn compile(self, mode: &compile::Mode, module_name: String) -> Result<CodeObject, Diagnostic> {
         let compile = |source| {
-            compile::compile(source, mode, source_path).map_err(|err| {
+            compile::compile(source, mode, module_name).map_err(|err| {
                 Diagnostic::spans_error(self.span, format!("Compile error: {}", err))
             })
         };
@@ -69,7 +69,7 @@ struct PyCompileInput {
 
 impl PyCompileInput {
     fn compile(&self) -> Result<CodeObject, Diagnostic> {
-        let mut source_path = None;
+        let mut module_name = None;
         let mut mode = None;
         let mut source: Option<CompilationSource> = None;
 
@@ -97,10 +97,10 @@ impl PyCompileInput {
                             },
                             _ => bail_span!(name_value.lit, "mode must be a string"),
                         })
-                    } else if name_value.ident == "source_path" {
-                        source_path = Some(match &name_value.lit {
+                    } else if name_value.ident == "module_name" {
+                        module_name = Some(match &name_value.lit {
                             Lit::Str(s) => s.value(),
-                            _ => bail_span!(name_value.lit, "source_path must be string"),
+                            _ => bail_span!(name_value.lit, "module_name must be string"),
                         })
                     } else if name_value.ident == "source" {
                         assert_source_empty(&source)?;
@@ -137,7 +137,7 @@ impl PyCompileInput {
             })?
             .compile(
                 &mode.unwrap_or(compile::Mode::Exec),
-                source_path.unwrap_or_else(|| "frozen".to_string()),
+                module_name.unwrap_or_else(|| "frozen".to_string()),
             )
     }
 }

vm/src/frozen.rs

@@ -5,12 +5,15 @@ pub fn get_module_inits() -> HashMap<String, CodeObject> {
     hashmap! {
         "__hello__".into() => py_compile_bytecode!(
             source = "initialized = True; print(\"Hello world!\")\n",
+            module_name = "__hello__",
         ),
         "_frozen_importlib".into() => py_compile_bytecode!(
             file = "../Lib/importlib/_bootstrap.py",
+            module_name = "_frozen_importlib",
         ),
         "_frozen_importlib_external".into() => py_compile_bytecode!(
             file = "../Lib/importlib/_bootstrap_external.py",
+            module_name = "_frozen_importlib_external",
         ),
     }
 }

vm/src/import.rs

@@ -25,24 +25,24 @@ pub fn init_importlib(vm: &VirtualMachine) -> PyResult {
 }
 
 pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult {
-    if let Some(frozen) = vm.frozen.borrow().get(module_name) {
-        let mut frozen = frozen.clone();
-        frozen.source_path = format!("frozen {}", module_name);
-        import_codeobj(vm, module_name, frozen)
-    } else {
-        Err(vm.new_import_error(format!("Cannot import frozen module {}", module_name)))
-    }
+    vm.frozen
+        .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))
 }
 
 pub fn import_builtin(vm: &VirtualMachine, module_name: &str) -> PyResult {
-    let sys_modules = vm.get_attribute(vm.sys_module.clone(), "modules").unwrap();
-    if let Some(make_module_func) = vm.stdlib_inits.borrow().get(module_name) {
-        let module = make_module_func(vm);
-        sys_modules.set_item(module_name, module.clone(), vm)?;
-        Ok(module)
-    } else {
-        Err(vm.new_import_error(format!("Cannot import bultin module {}", module_name)))
-    }
+    vm.stdlib_inits
+        .borrow()
+        .get(module_name)
+        .ok_or_else(|| vm.new_import_error(format!("Cannot import bultin module {}", module_name)))
+        .and_then(|make_module_func| {
+            let module = make_module_func(vm);
+            let sys_modules = vm.get_attribute(vm.sys_module.clone(), "modules")?;
+            sys_modules.set_item(module_name, module.clone(), vm)?;
+            Ok(module)
+        })
 }
 
 pub fn import_module(vm: &VirtualMachine, current_path: PathBuf, module_name: &str) -> PyResult {
@@ -57,8 +57,8 @@ pub fn import_module(vm: &VirtualMachine, current_path: PathBuf, module_name: &s
     } else if vm.stdlib_inits.borrow().contains_key(module_name) {
         import_builtin(vm, module_name)
     } else {
-        let notfound_error = vm.context().exceptions.module_not_found_error.clone();
-        let import_error = vm.context().exceptions.import_error.clone();
+        let notfound_error = &vm.ctx.exceptions.module_not_found_error;
+        let import_error = &vm.ctx.exceptions.import_error;
 
         // Time to search for module in any place:
         let file_path = find_source(vm, current_path, module_name)
@@ -83,21 +83,24 @@ pub fn import_file(
 ) -> PyResult {
     let code_obj = compile::compile(&content, &compile::Mode::Exec, file_path)
         .map_err(|err| vm.new_syntax_error(&err))?;
-    import_codeobj(vm, module_name, code_obj)
+    import_codeobj(vm, module_name, code_obj, true)
 }
 
-pub fn import_codeobj(vm: &VirtualMachine, module_name: &str, code_obj: CodeObject) -> PyResult {
+pub fn import_codeobj(
+    vm: &VirtualMachine,
+    module_name: &str,
+    code_obj: CodeObject,
+    set_file_attr: bool,
+) -> PyResult {
     let attrs = vm.ctx.new_dict();
     attrs.set_item("__name__", vm.new_str(module_name.to_string()), vm)?;
-    let file_path = &code_obj.source_path;
-    if !file_path.starts_with("frozen") {
-        // TODO: Should be less hacky, not depend on source_path
-        attrs.set_item("__file__", vm.new_str(file_path.to_owned()), vm)?;
+    if set_file_attr {
+        attrs.set_item("__file__", vm.new_str(code_obj.source_path.to_owned()), vm)?;
     }
     let module = vm.ctx.new_module(module_name, attrs.clone());
 
     // Store module in cache to prevent infinite loop with mutual importing libs:
-    let sys_modules = vm.get_attribute(vm.sys_module.clone(), "modules").unwrap();
+    let sys_modules = vm.get_attribute(vm.sys_module.clone(), "modules")?;
     sys_modules.set_item(module_name, module.clone(), vm)?;
 
     // Execute main code in module: