From 090cfed5930067f729dcfe2c27b5bdfc9753e6a4 Mon Sep 17 00:00:00 2001
From: Aviv Palivoda <palaviv@gmail.com>
Date: Wed, 3 Jul 2019 18:57:49 +0300
Subject: [PATCH] Use importlib import in WASM

---
 src/main.rs                   |  2 +-
 vm/src/import.rs              |  9 ++++++---
 wasm/lib/src/vm_class.rs      |  3 ++-
 wasm/lib/src/wasm_builtins.rs | 27 ---------------------------
 4 files changed, 9 insertions(+), 32 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index ba6f097ec9a..c9f74fea0a0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -51,7 +51,7 @@ fn main() {
     // Construct vm:
     let vm = VirtualMachine::new();
 
-    let res = import::init_importlib(&vm);
+    let res = import::init_importlib(&vm, true);
     handle_exception(&vm, res);
 
     // Figure out if a -c option was given:
diff --git a/vm/src/import.rs b/vm/src/import.rs
index 10d2175054f..37fcca3385a 100644
--- a/vm/src/import.rs
+++ b/vm/src/import.rs
@@ -13,15 +13,18 @@ use crate::vm::VirtualMachine;
 #[cfg(feature = "rustpython-compiler")]
 use rustpython_compiler::compile;
 
-pub fn init_importlib(vm: &VirtualMachine) -> PyResult {
+pub fn init_importlib(vm: &VirtualMachine, external: bool) -> 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![])?;
+    if external {
+        let install_external =
+            vm.get_attribute(importlib.clone(), "_install_external_importers")?;
+        vm.invoke(install_external, vec![])?;
+    }
     Ok(vm.get_none())
 }
 
diff --git a/wasm/lib/src/vm_class.rs b/wasm/lib/src/vm_class.rs
index 2bb156ef8c2..1437af46683 100644
--- a/wasm/lib/src/vm_class.rs
+++ b/wasm/lib/src/vm_class.rs
@@ -8,6 +8,7 @@ use wasm_bindgen::prelude::*;
 use rustpython_compiler::{compile, error::CompileErrorType};
 use rustpython_vm::frame::{NameProtocol, Scope};
 use rustpython_vm::function::PyFuncArgs;
+use rustpython_vm::import;
 use rustpython_vm::pyobject::{PyObject, PyObjectPayload, PyObjectRef, PyResult, PyValue};
 use rustpython_vm::VirtualMachine;
 
@@ -43,7 +44,7 @@ impl StoredVirtualMachine {
             setup_browser_module(&vm);
         }
 
-        *vm.import_func.borrow_mut() = vm.ctx.new_rustfunc(wasm_builtins::builtin_import);
+        import::init_importlib(&vm, false);
 
         StoredVirtualMachine {
             vm,
diff --git a/wasm/lib/src/wasm_builtins.rs b/wasm/lib/src/wasm_builtins.rs
index 6a89f5e35f0..e2fed62c2bc 100644
--- a/wasm/lib/src/wasm_builtins.rs
+++ b/wasm/lib/src/wasm_builtins.rs
@@ -8,7 +8,6 @@ use js_sys::{self, Array};
 use web_sys::{self, console};
 
 use rustpython_vm::function::{Args, KwArgs, PyFuncArgs};
-use rustpython_vm::import;
 use rustpython_vm::obj::{
     objstr::{self, PyStringRef},
     objtype,
@@ -80,29 +79,3 @@ pub fn builtin_print_console(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult
     console::log(&arr);
     Ok(vm.get_none())
 }
-
-pub fn builtin_import(
-    module_name: PyStringRef,
-    _args: Args,
-    _kwargs: KwArgs,
-    vm: &VirtualMachine,
-) -> PyResult<PyObjectRef> {
-    let module_name = module_name.as_str();
-
-    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 vm.frozen.borrow().contains_key(module_name) {
-        import::import_frozen(vm, module_name)
-    } else if vm.stdlib_inits.borrow().contains_key(module_name) {
-        import::import_builtin(vm, module_name)
-    } else {
-        let notfound_error = vm.context().exceptions.module_not_found_error.clone();
-        Err(vm.new_exception(
-            notfound_error,
-            format!("Module {:?} not found", module_name),
-        ))
-    }
-}