windelbouwman · Jun 12, 2019 · Jun 1, 2019 · Jun 5, 2019 · Jun 5, 2019 · Jun 7, 2019
diff --git a/Lib/importlib/_bootstrap.py b/Lib/importlib/_bootstrap.py
diff --git a/Lib/importlib/_bootstrap_external.py b/Lib/importlib/_bootstrap_external.py
diff --git a/Lib/io.py b/Lib/io.py
@@ -0,0 +1 @@
+from _io import *
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())
@@ -126,6 +129,10 @@ fn run_script(vm: &VirtualMachine, script_file: &str) -> PyResult {
         std::process::exit(1);
     };
 
+    let dir = file_path.parent().unwrap().to_str().unwrap().to_string();
+    let sys_path = vm.get_attribute(vm.sys_module.clone(), "path").unwrap();
+    vm.call_method(&sys_path, "insert", vec![vm.new_int(0), vm.new_str(dir)])?;
+
     match util::read_file(&file_path) {
         Ok(source) => _run_string(vm, &source, file_path.to_str().unwrap().to_string()),
         Err(err) => {

diff --git a/tests/snippets/import.py b/tests/snippets/import.py
@@ -7,6 +7,8 @@
 assert import_target.X == import_target.func()
 assert import_target.X == func()
 
+assert import_mutual1.__name__ == "import_mutual1"
+
 assert import_target.Y == other_func()
 
 assert import_target.X == aliased.X

diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -4,14 +4,14 @@
 
 use std::char;
 use std::io::{self, Write};
-use std::path::PathBuf;
+use std::str;
 
 use num_bigint::Sign;
 use num_traits::{Signed, Zero};
 
 use crate::compile;
-use crate::import::import_module;
 use crate::obj::objbool;
+use crate::obj::objbytes::PyBytesRef;
 use crate::obj::objcode::PyCodeRef;
 use crate::obj::objdict::PyDictRef;
 use crate::obj::objint::{self, PyIntRef};
@@ -22,7 +22,7 @@ use crate::obj::objtype::{self, PyClassRef};
 use crate::frame::Scope;
 use crate::function::{single_or_tuple_any, Args, KwArgs, OptionalArg, PyFuncArgs};
 use crate::pyobject::{
-    IdProtocol, IntoPyObject, ItemProtocol, PyIterable, PyObjectRef, PyResult, PyValue,
+    Either, IdProtocol, IntoPyObject, ItemProtocol, PyIterable, PyObjectRef, PyResult, PyValue,
     TryFromObject, TypeProtocol,
 };
 use crate::vm::VirtualMachine;
@@ -79,17 +79,35 @@ fn builtin_chr(i: u32, vm: &VirtualMachine) -> PyResult<String> {
     }
 }
 
-fn builtin_compile(
-    source: PyStringRef,
+#[derive(FromArgs)]
+#[allow(dead_code)]
+struct CompileArgs {
+    #[pyarg(positional_only, optional = false)]
+    source: Either<PyStringRef, PyBytesRef>,
+    #[pyarg(positional_only, optional = false)]
     filename: PyStringRef,
+    #[pyarg(positional_only, optional = false)]
     mode: PyStringRef,
-    vm: &VirtualMachine,
-) -> PyResult<PyCodeRef> {
+    #[pyarg(positional_or_keyword, optional = true)]
+    flags: OptionalArg<PyIntRef>,
+    #[pyarg(positional_or_keyword, optional = true)]
+    dont_inherit: OptionalArg<bool>,
+    #[pyarg(positional_or_keyword, optional = true)]
+    optimize: OptionalArg<PyIntRef>,
+}
+
+fn builtin_compile(args: CompileArgs, vm: &VirtualMachine) -> PyResult<PyCodeRef> {
+    // TODO: compile::compile should probably get bytes
+    let source = match args.source {
+        Either::A(string) => string.value.to_string(),
+        Either::B(bytes) => str::from_utf8(&bytes).unwrap().to_string(),
+    };
+
     // TODO: fix this newline bug:
-    let source = format!("{}\n", &source.value);
+    let source = format!("{}\n", source);
 
     let mode = {
-        let mode = &mode.value;
+        let mode = &args.mode.value;
         if mode == "exec" {
             compile::Mode::Exec
         } else if mode == "eval" {
@@ -103,7 +121,7 @@ fn builtin_compile(
         }
     };
 
-    compile::compile(vm, &source, &mode, filename.value.to_string())
+    compile::compile(vm, &source, &mode, args.filename.value.to_string())
         .map_err(|err| vm.new_syntax_error(&err))
 }
 
@@ -732,27 +750,7 @@ fn builtin_sum(iterable: PyIterable, start: OptionalArg, vm: &VirtualMachine) ->
 
 // Should be renamed to builtin___import__?
 fn builtin_import(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(
-        vm,
-        args,
-        required = [(name, Some(vm.ctx.str_type()))],
-        optional = [
-            (_globals, Some(vm.ctx.dict_type())),
-            (_locals, Some(vm.ctx.dict_type()))
-        ]
-    );
-    let current_path = {
-        match vm.current_frame() {
-            Some(frame) => {
-                let mut source_pathbuf = PathBuf::from(&frame.code.source_path);
-                source_pathbuf.pop();
-                source_pathbuf
-            }
-            None => PathBuf::new(),
-        }
-    };
-
-    import_module(vm, current_path, &objstr::get_value(name))
+    vm.invoke(vm.import_func.borrow().clone(), args)
 }
 
 // builtin_vars
@@ -860,6 +858,7 @@ pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) {
         "ZeroDivisionError" => ctx.exceptions.zero_division_error.clone(),
         "KeyError" => ctx.exceptions.key_error.clone(),
         "OSError" => ctx.exceptions.os_error.clone(),
+        "ModuleNotFoundError" => ctx.exceptions.module_not_found_error.clone(),
 
         // Warnings
         "Warning" => ctx.exceptions.warning.clone(),

diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -908,7 +908,11 @@ impl Frame {
     }
 
     fn import(&self, vm: &VirtualMachine, module: &str, symbol: &Option<String>) -> FrameResult {
-        let module = vm.import(module)?;
+        let from_list = match symbol {
+            Some(symbol) => vm.ctx.new_tuple(vec![vm.ctx.new_str(symbol.to_string())]),
+            None => vm.ctx.new_tuple(vec![]),
+        };
+        let module = vm.import(module, &from_list)?;
 
         // If we're importing a symbol, look it up and use it, otherwise construct a module and return
         // that
@@ -926,7 +930,7 @@ impl Frame {
     }
 
     fn import_star(&self, vm: &VirtualMachine, module: &str) -> FrameResult {
-        let module = vm.import(module)?;
+        let module = vm.import(module, &vm.ctx.new_tuple(vec![]))?;
 
         // Grab all the names from the module and put them in the context
         if let Some(dict) = &module.dict {

diff --git a/vm/src/frozen.rs b/vm/src/frozen.rs
@@ -4,8 +4,22 @@ const HELLO: &str = "initialized = True
 print(\"Hello world!\")
 ";
 
+const IMPORTLIB_BOOTSTRAP: &'static str = include_str!(concat!(
+    env!("CARGO_MANIFEST_DIR"),
+    "/../Lib/importlib/_bootstrap.py"
+));
+const IMPORTLIB_BOOTSTRAP_EXTERNAL: &'static str = include_str!(concat!(
+    env!("CARGO_MANIFEST_DIR"),
+    "/../Lib/importlib/_bootstrap_external.py"
+));
+
 pub fn get_module_inits() -> HashMap<String, &'static str> {
     let mut modules = HashMap::new();
     modules.insert("__hello__".to_string(), HELLO);
+    modules.insert("_frozen_importlib".to_string(), IMPORTLIB_BOOTSTRAP);
+    modules.insert(
+        "_frozen_importlib_external".to_string(),
+        IMPORTLIB_BOOTSTRAP_EXTERNAL,
+    );
     modules
 }
diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -11,19 +11,53 @@ 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.clone(), "_install")?;
+    vm.invoke(install, vec![vm.sys_module.clone(), impmod])?;
+    vm.import_func
+        .replace(vm.get_attribute(importlib.clone(), "__import__")?);
+    let install_external = vm.get_attribute(importlib.clone(), "_install_external_importers")?;
+    vm.invoke(install_external, vec![])?;
+    Ok(vm.get_none())
+}
+
+pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult {
+    if let Some(frozen) = vm.frozen.borrow().get(module_name) {
+        import_file(
+            vm,
+            module_name,
+            format!("frozen {}", module_name),
+            frozen.to_string(),
+        )
+    } else {
+        Err(vm.new_import_error(format!("Cannot import frozen module {}", module_name)))
+    }
+}
+
+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)))
+    }
+}
+
 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();
@@ -56,7 +90,10 @@ pub fn import_file(
 
     let attrs = vm.ctx.new_dict();
     attrs.set_item("__name__", vm.new_str(module_name.to_string()), vm)?;
-    attrs.set_item("__file__", vm.new_str(file_path), vm)?;
+    if !file_path.starts_with("frozen") {
+        // TODO: Should be removed after precompiling frozen modules.
+        attrs.set_item("__file__", vm.new_str(file_path), vm)?;
+    }
     let module = vm.ctx.new_module(module_name, attrs.clone());
 
     // Store module in cache to prevent infinite loop with mutual importing libs:

diff --git a/vm/src/macros.rs b/vm/src/macros.rs
@@ -118,6 +118,7 @@ macro_rules! no_kwargs {
 macro_rules! py_module {
     ( $vm:expr, $module_name:expr, { $($name:expr => $value:expr),* $(,)* }) => {{
         let module = $vm.ctx.new_module($module_name, $vm.ctx.new_dict());
+        $vm.set_attr(&module, "__name__", $vm.ctx.new_str($module_name.to_string())).unwrap();
         $(
             $vm.set_attr(&module, $name, $value).unwrap();
         )*

diff --git a/vm/src/obj/objmodule.rs b/vm/src/obj/objmodule.rs
@@ -1,3 +1,4 @@
+use crate::obj::objstr::PyStringRef;
 use crate::obj::objtype::PyClassRef;
 use crate::pyobject::{PyContext, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
@@ -9,24 +10,22 @@ pub struct PyModule {
 pub type PyModuleRef = PyRef<PyModule>;
 
 impl PyValue for PyModule {
+    const HAVE_DICT: bool = true;
+
     fn class(vm: &VirtualMachine) -> PyClassRef {
         vm.ctx.module_type()
     }
 }
 
 impl PyModuleRef {
-    fn dir(self: PyModuleRef, vm: &VirtualMachine) -> PyResult {
-        if let Some(dict) = &self.into_object().dict {
-            let keys = dict.into_iter().map(|(k, _v)| k.clone()).collect();
-            Ok(vm.ctx.new_list(keys))
-        } else {
-            panic!("Modules should definitely have a dict.");
-        }
+    fn init(self, name: PyStringRef, vm: &VirtualMachine) -> PyResult {
+        vm.set_attr(&self.into_object(), "__name__", name)?;
+        Ok(vm.get_none())
     }
 }
 
 pub fn init(context: &PyContext) {
     extend_class!(&context, &context.module_type, {
-        "__dir__" => context.new_rustfunc(PyModuleRef::dir),
+        "__init__" => context.new_rustfunc(PyModuleRef::init),
     });
 }
diff --git a/vm/src/stdlib/imp.rs b/vm/src/stdlib/imp.rs
@@ -1,5 +1,5 @@
 use crate::compile;
-use crate::import::import_file;
+use crate::import;
 use crate::obj::objcode::PyCodeRef;
 use crate::obj::objmodule::PyModuleRef;
 use crate::obj::objstr;
@@ -56,20 +56,22 @@ fn imp_exec_builtin(_mod: PyModuleRef, _vm: &VirtualMachine) -> i32 {
 }
 
 fn imp_get_frozen_object(name: PyStringRef, vm: &VirtualMachine) -> PyResult<PyCodeRef> {
-    if let Some(frozen) = vm.frozen.borrow().get(name.as_str()) {
-        compile::compile(vm, frozen, &compile::Mode::Exec, "frozen".to_string())
-            .map_err(|err| vm.new_syntax_error(&err))
+    let name_str = name.as_str();
+    if let Some(frozen) = vm.frozen.borrow().get(name_str) {
+        compile::compile(
+            vm,
+            frozen,
+            &compile::Mode::Exec,
+            format!("frozen {}", name_str),
+        )
+        .map_err(|err| vm.new_syntax_error(&err))
     } else {
         Err(vm.new_import_error(format!("No such frozen object named {}", name.as_str())))
     }
 }
 
 fn imp_init_frozen(name: PyStringRef, vm: &VirtualMachine) -> PyResult {
-    if let Some(frozen) = vm.frozen.borrow().get(name.as_str()) {
-        import_file(vm, name.as_str(), "frozen".to_string(), frozen.to_string())
-    } else {
-        Err(vm.new_import_error(format!("No such frozen object named {}", name.as_str())))
-    }
+    import::import_frozen(vm, name.as_str())
 }
 
 fn imp_is_frozen_package(_name: PyStringRef, _vm: &VirtualMachine) -> bool {

diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs
@@ -471,7 +471,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "getvalue" => ctx.new_rustfunc(bytes_io_getvalue)
     });
 
-    py_module!(vm, "io", {
+    py_module!(vm, "_io", {
         "open" => ctx.new_rustfunc(io_open),
         "IOBase" => io_base,
         "RawIOBase" => raw_io_base,

diff --git a/vm/src/stdlib/mod.rs b/vm/src/stdlib/mod.rs
@@ -61,7 +61,7 @@ pub fn get_module_inits() -> HashMap<String, StdlibInitFunc> {
     // disable some modules on WASM
     #[cfg(not(target_arch = "wasm32"))]
     {
-        modules.insert("io".to_string(), Box::new(io::make_module));
+        modules.insert("_io".to_string(), Box::new(io::make_module));
         modules.insert("_os".to_string(), Box::new(os::make_module));
         modules.insert("socket".to_string(), Box::new(socket::make_module));
     }

diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs
@@ -234,7 +234,7 @@ type DirEntryRef = PyRef<DirEntry>;
 
 impl PyValue for DirEntry {
     fn class(vm: &VirtualMachine) -> PyClassRef {
-        vm.class("os", "DirEntry")
+        vm.class("_os", "DirEntry")
     }
 }
 
@@ -316,7 +316,7 @@ struct ScandirIterator {
 
 impl PyValue for ScandirIterator {
     fn class(vm: &VirtualMachine) -> PyClassRef {
-        vm.class("os", "ScandirIter")
+        vm.class("_os", "ScandirIter")
     }
 }
 
@@ -366,7 +366,7 @@ struct StatResult {
 
 impl PyValue for StatResult {
     fn class(vm: &VirtualMachine) -> PyClassRef {
-        vm.class("os", "stat_result")
+        vm.class("_os", "stat_result")
     }
 }