Add init_importlib

palaviv · palaviv · commit e88d6ac1ef5f · 2019-06-10T19:18:32.000+03:00
diff --git a/src/main.rs b/src/main.rs
@@ -54,6 +54,9 @@ fn main() {
     // Construct vm:
     let vm = VirtualMachine::new();
 
+    let res = import::init_importlib(&vm);
+    handle_exception(&vm, res);
+
     // Figure out if a -c option was given:
     let result = if let Some(command) = matches.value_of("c") {
         run_command(&vm, command.to_string())
diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -11,19 +11,43 @@ use crate::pyobject::{ItemProtocol, PyResult};
 use crate::util;
 use crate::vm::VirtualMachine;
 
+pub fn init_importlib(vm: &VirtualMachine) -> PyResult {
+    let importlib = import_frozen(vm, "_frozen_importlib")?;
+    let impmod = import_builtin(vm, "_imp")?;
+    let install = vm.get_attribute(importlib, "_install")?;
+    vm.invoke(install, vec![vm.sys_module.clone(), impmod])
+}
+
+fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult {
+    if let Some(frozen) = vm.frozen.borrow().get(module_name) {
+        import_file(vm, module_name, "frozen".to_string(), frozen.to_string())
+    } else {
+        Err(vm.new_import_error(format!("Cannot import frozen module {}", module_name)))
+    }
+}
+
+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)))
+    }
+}
+
 pub fn import_module(vm: &VirtualMachine, current_path: PathBuf, module_name: &str) -> PyResult {
     // Cached modules:
     let sys_modules = vm.get_attribute(vm.sys_module.clone(), "modules").unwrap();
 
     // First, see if we already loaded the module:
     if let Ok(module) = sys_modules.get_item(module_name.to_string(), vm) {
         Ok(module)
-    } else if let Some(frozen) = vm.frozen.borrow().get(module_name) {
-        import_file(vm, module_name, "frozen".to_string(), frozen.to_string())
-    } else 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 if vm.frozen.borrow().contains_key(module_name) {
+        import_frozen(vm, module_name)
+    } 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();
