Merge pull request RustPython#1011 from palaviv/_imp2

coolreader18 · web-flow · commit 74d52ee82f65 · 2019-06-05T00:43:10.000-05:00
More _imp
diff --git a/tests/snippets/imp.py b/tests/snippets/imp.py
@@ -1,5 +1,28 @@
 import _imp
+import time as import_time
 
 assert _imp.is_builtin("time") == True
 assert _imp.is_builtin("os") == False
 assert _imp.is_builtin("not existing module") == False
+
+assert _imp.is_frozen("__hello__") == True
+assert _imp.is_frozen("os") == False
+
+class FakeSpec:
+	def __init__(self, name):
+		self.name = name
+
+A = FakeSpec("time")
+
+imp_time = _imp.create_builtin(A)
+assert imp_time.sleep == import_time.sleep
+
+B = FakeSpec("not existing module")
+assert _imp.create_builtin(B) == None
+
+_imp.exec_builtin(imp_time) == 0
+
+_imp.get_frozen_object("__hello__")
+
+hello = _imp.init_frozen("__hello__")
+assert hello.initialized == True
diff --git a/vm/src/stdlib/imp.rs b/vm/src/stdlib/imp.rs
@@ -1,5 +1,10 @@
+use crate::compile;
+use crate::import::import_file;
+use crate::obj::objcode::PyCodeRef;
+use crate::obj::objmodule::PyModuleRef;
+use crate::obj::objstr;
 use crate::obj::objstr::PyStringRef;
-use crate::pyobject::{PyObjectRef, PyResult};
+use crate::pyobject::{ItemProtocol, PyObjectRef, PyResult};
 use crate::vm::VirtualMachine;
 
 fn imp_extension_suffixes(vm: &VirtualMachine) -> PyResult {
@@ -25,6 +30,53 @@ fn imp_is_builtin(name: PyStringRef, vm: &VirtualMachine) -> bool {
     vm.stdlib_inits.borrow().contains_key(name.as_str())
 }
 
+fn imp_is_frozen(name: PyStringRef, vm: &VirtualMachine) -> bool {
+    vm.frozen.borrow().contains_key(name.as_str())
+}
+
+fn imp_create_builtin(spec: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+    let sys_modules = vm.get_attribute(vm.sys_module.clone(), "modules").unwrap();
+
+    let name = &objstr::get_value(&vm.get_attribute(spec.clone(), "name")?);
+
+    if let Ok(module) = sys_modules.get_item(name, vm) {
+        Ok(module)
+    } else {
+        if let Some(make_module_func) = vm.stdlib_inits.borrow().get(name) {
+            Ok(make_module_func(vm))
+        } else {
+            Ok(vm.get_none())
+        }
+    }
+}
+
+fn imp_exec_builtin(_mod: PyModuleRef, _vm: &VirtualMachine) -> i32 {
+    // TOOD: Should we do something here?
+    0
+}
+
+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))
+    } 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())))
+    }
+}
+
+fn imp_is_frozen_package(_name: PyStringRef, _vm: &VirtualMachine) -> bool {
+    // TODO: Support frozen package.
+    false
+}
+
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let ctx = &vm.ctx;
     let module = py_module!(vm, "_imp", {
@@ -33,6 +85,12 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "release_lock" => ctx.new_rustfunc(imp_release_lock),
         "lock_held" => ctx.new_rustfunc(imp_lock_held),
         "is_builtin" => ctx.new_rustfunc(imp_is_builtin),
+        "is_frozen" => ctx.new_rustfunc(imp_is_frozen),
+        "create_builtin" => ctx.new_rustfunc(imp_create_builtin),
+        "exec_builtin" => ctx.new_rustfunc(imp_exec_builtin),
+        "get_frozen_object" => ctx.new_rustfunc(imp_get_frozen_object),
+        "init_frozen" => ctx.new_rustfunc(imp_init_frozen),
+        "is_frozen_package" => ctx.new_rustfunc(imp_is_frozen_package),
     });
 
     module
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -255,6 +255,11 @@ impl VirtualMachine {
         syntax_error
     }
 
+    pub fn new_import_error(&self, msg: String) -> PyObjectRef {
+        let import_error = self.ctx.exceptions.import_error.clone();
+        self.new_exception(import_error, msg)
+    }
+
     pub fn new_scope_with_builtins(&self) -> Scope {
         Scope::with_builtins(None, self.ctx.new_dict(), self)
     }

Original file line number	Diff line number	Diff line change
`@@ -255,6 +255,11 @@ impl VirtualMachine {`
`255`	`255`	`syntax_error`
`256`	`256`	`}`
`257`	`257`
	`258`	`+ pub fn new_import_error(&self, msg: String) -> PyObjectRef {`
	`259`	`+ let import_error = self.ctx.exceptions.import_error.clone();`
	`260`	`+ self.new_exception(import_error, msg)`
	`261`	`+ }`
	`262`	`+`
`258`	`263`	`pub fn new_scope_with_builtins(&self) -> Scope {`
`259`	`264`	`Scope::with_builtins(None, self.ctx.new_dict(), self)`
`260`	`265`	`}`