Merge pull request #1018 from palaviv/imp3

Importlib

Lib/io.py

@@ -0,0 +1 @@
+from _io import *

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) => {

tests/snippets/import.py

@@ -7,6 +7,8 @@ import import_mutual1
 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

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))
 }
 
@@ -724,27 +742,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
@@ -852,6 +850,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(),

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 {

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
 }

vm/src/import.rs

@@ -11,6 +11,42 @@ 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();
@@ -18,12 +54,10 @@ pub fn import_module(vm: &VirtualMachine, current_path: PathBuf, module_name: &s
     // 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:

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();
         )*

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),
     });
 }

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 {

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,

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));
     }

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")
     }
 }
 

vm/src/stdlib/thread.rs

@@ -2,8 +2,10 @@
 /// support threading
 use super::super::pyobject::PyObjectRef;
 use crate::function::PyFuncArgs;
+use crate::import;
 use crate::pyobject::PyResult;
 use crate::vm::VirtualMachine;
+use std::path::PathBuf;
 
 fn rlock_acquire(vm: &VirtualMachine, _args: PyFuncArgs) -> PyResult {
     Ok(vm.get_none())
@@ -11,20 +13,49 @@ fn rlock_acquire(vm: &VirtualMachine, _args: PyFuncArgs) -> PyResult {
 
 fn rlock_release(_zelf: PyObjectRef, _vm: &VirtualMachine) {}
 
+fn rlock_enter(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(instance, None)]);
+    Ok(instance.clone())
+}
+
+fn rlock_exit(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        // The context manager protocol requires these, but we don't use them
+        required = [
+            (_instance, None),
+            (_exception_type, None),
+            (_exception_value, None),
+            (_traceback, None)
+        ]
+    );
+    Ok(vm.get_none())
+}
+
 fn get_ident(_vm: &VirtualMachine) -> u32 {
     1
 }
 
+fn allocate_lock(vm: &VirtualMachine) -> PyResult {
+    let module = import::import_module(vm, PathBuf::default(), "_thread")?;
+    let lock_class = vm.get_attribute(module.clone(), "RLock")?;
+    vm.invoke(lock_class, vec![])
+}
+
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let ctx = &vm.ctx;
 
     let rlock_type = py_class!(ctx, "_thread.RLock", ctx.object(), {
         "acquire" => ctx.new_rustfunc(rlock_acquire),
         "release" => ctx.new_rustfunc(rlock_release),
+        "__enter__" => ctx.new_rustfunc(rlock_enter),
+        "__exit__" => ctx.new_rustfunc(rlock_exit),
     });
 
     py_module!(vm, "_thread", {
         "RLock" => rlock_type,
-        "get_ident" => ctx.new_rustfunc(get_ident)
+        "get_ident" => ctx.new_rustfunc(get_ident),
+        "allocate_lock" => ctx.new_rustfunc(allocate_lock),
     })
 }

vm/src/sysmodule.rs

@@ -224,6 +224,7 @@ settrace() -- set the global debug tracing function
     module_names.sort();
     let modules = ctx.new_dict();
     extend_module!(vm, module, {
+      "__name__" => ctx.new_str(String::from("sys")),
       "argv" => argv(ctx),
       "builtin_module_names" => ctx.new_tuple(module_names.iter().map(|v| v.into_pyobject(vm).unwrap()).collect()),
       "flags" => flags,
@@ -242,6 +243,11 @@ settrace() -- set the global debug tracing function
       "modules" => modules.clone(),
       "warnoptions" => ctx.new_list(vec![]),
       "platform" => ctx.new_str(platform),
+      "meta_path" => ctx.new_list(vec![]),
+      "path_hooks" => ctx.new_list(vec![]),
+      "path_importer_cache" => ctx.new_dict(),
+      "pycache_prefix" => vm.get_none(),
+      "dont_write_bytecode" => vm.new_bool(true),
     });
 
     modules.set_item("sys", module.clone(), vm).unwrap();

vm/src/vm.rs

@@ -55,6 +55,7 @@ pub struct VirtualMachine {
     pub wasm_id: Option<String>,
     pub exceptions: RefCell<Vec<PyObjectRef>>,
     pub frozen: RefCell<HashMap<String, &'static str>>,
+    pub import_func: RefCell<PyObjectRef>,
 }
 
 impl VirtualMachine {
@@ -68,6 +69,7 @@ impl VirtualMachine {
 
         let stdlib_inits = RefCell::new(stdlib::get_module_inits());
         let frozen = RefCell::new(frozen::get_module_inits());
+        let import_func = RefCell::new(ctx.none());
         let vm = VirtualMachine {
             builtins: builtins.clone(),
             sys_module: sysmod.clone(),
@@ -77,6 +79,7 @@ impl VirtualMachine {
             wasm_id: None,
             exceptions: RefCell::new(vec![]),
             frozen,
+            import_func,
         };
 
         builtins::make_module(&vm, builtins.clone());
@@ -134,7 +137,7 @@ impl VirtualMachine {
 
     pub fn try_class(&self, module: &str, class: &str) -> PyResult<PyClassRef> {
         let class = self
-            .get_attribute(self.import(module)?, class)?
+            .get_attribute(self.import(module, &self.ctx.new_tuple(vec![]))?, class)?
             .downcast()
             .expect("not a class");
         Ok(class)
@@ -142,7 +145,7 @@ impl VirtualMachine {
 
     pub fn class(&self, module: &str, class: &str) -> PyClassRef {
         let module = self
-            .import(module)
+            .import(module, &self.ctx.new_tuple(vec![]))
             .unwrap_or_else(|_| panic!("unable to import {}", module));
         let class = self
             .get_attribute(module.clone(), class)
@@ -300,13 +303,28 @@ impl VirtualMachine {
         TryFromObject::try_from_object(self, repr)
     }
 
-    pub fn import(&self, module: &str) -> PyResult {
-        match self.get_attribute(self.builtins.clone(), "__import__") {
-            Ok(func) => self.invoke(func, vec![self.ctx.new_str(module.to_string())]),
-            Err(_) => Err(self.new_exception(
-                self.ctx.exceptions.import_error.clone(),
-                "__import__ not found".to_string(),
-            )),
+    pub fn import(&self, module: &str, from_list: &PyObjectRef) -> PyResult {
+        let sys_modules = self
+            .get_attribute(self.sys_module.clone(), "modules")
+            .unwrap();
+        if let Ok(module) = sys_modules.get_item(module.to_string(), self) {
+            Ok(module)
+        } else {
+            match self.get_attribute(self.builtins.clone(), "__import__") {
+                Ok(func) => self.invoke(
+                    func,
+                    vec![
+                        self.ctx.new_str(module.to_string()),
+                        self.get_none(),
+                        self.get_none(),
+                        from_list.clone(),
+                    ],
+                ),
+                Err(_) => Err(self.new_exception(
+                    self.ctx.exceptions.import_error.clone(),
+                    "__import__ not found".to_string(),
+                )),
+            }
         }
     }