Allow injecting JS variables into python with eval_py()

coolreader18 · coolreader18 · commit a796b134a5ba · 2018-12-15T12:20:03.000-06:00
eval_py(`return js_vars["a"]`, { a: 9 }) == 9
diff --git a/wasm/src/lib.rs b/wasm/src/lib.rs
@@ -5,6 +5,7 @@ extern crate rustpython_vm;
 extern crate wasm_bindgen;
 extern crate web_sys;
 
+use js_sys::Reflect;
 use rustpython_vm::compile;
 use rustpython_vm::pyobject::{self, PyObjectRef, PyResult};
 use rustpython_vm::VirtualMachine;
@@ -33,7 +34,31 @@ fn py_to_js(vm: &mut VirtualMachine, py_obj: PyObjectRef) -> JsValue {
     }
 }
 
-fn eval(vm: &mut VirtualMachine, source: &str) -> PyResult {
+fn js_to_py(vm: &mut VirtualMachine, js_val: JsValue) -> PyObjectRef {
+    let json = match js_sys::JSON::stringify(&js_val) {
+        Ok(json) => String::from(json),
+        Err(_) => return vm.get_none(),
+    };
+
+    let loads = rustpython_vm::import::import(
+        vm,
+        std::path::PathBuf::default(),
+        "json",
+        &Some("loads".into()),
+    )
+    .expect("Couldn't get json.loads function");
+
+    let py_json = vm.new_str(json);
+
+    vm.invoke(loads, pyobject::PyFuncArgs::new(vec![py_json], vec![]))
+        // can safely unwrap because we know it's valid JSON
+        .unwrap()
+}
+
+fn eval<F>(vm: &mut VirtualMachine, source: &str, setup_scope: F) -> PyResult
+where
+    F: Fn(&mut VirtualMachine, &PyObjectRef),
+{
     // HACK: if the code doesn't end with newline it crashes.
     let mut source = source.to_string();
     if !source.ends_with('\n') {
@@ -43,13 +68,15 @@ fn eval(vm: &mut VirtualMachine, source: &str) -> PyResult {
     let code_obj = compile::compile(vm, &source, compile::Mode::Exec, None)?;
 
     let builtins = vm.get_builtin_scope();
-    let vars = vm.context().new_scope(Some(builtins));
+    let mut vars = vm.context().new_scope(Some(builtins));
+
+    setup_scope(vm, &mut vars);
 
     vm.run_code_obj(code_obj, vars)
 }
 
 #[wasm_bindgen]
-pub fn eval_py(source: &str) -> Result<JsValue, JsValue> {
+pub fn eval_py(source: &str, js_injections: Option<js_sys::Object>) -> Result<JsValue, JsValue> {
     let mut vm = VirtualMachine::new();
 
     vm.ctx.set_attr(
@@ -59,8 +86,30 @@ pub fn eval_py(source: &str) -> Result<JsValue, JsValue> {
             .new_rustfunc(wasm_builtins::builtin_print_console),
     );
 
-    eval(&mut vm, source)
-        .map(|value| py_to_js(&mut vm, value))
+    let res = eval(&mut vm, source, |vm, vars| {
+        let injections = vm.new_dict();
+
+        if let Some(js_injections) = js_injections.clone() {
+            for pair in js_sys::try_iter(&js_sys::Object::entries(&js_injections))
+                .unwrap()
+                .unwrap()
+            {
+                let pair = pair.unwrap();
+                let key = Reflect::get(&pair, &"0".into()).unwrap();
+                let val = Reflect::get(&pair, &"1".into()).unwrap();
+                let py_val = js_to_py(vm, val);
+                vm.ctx.set_item(
+                    &injections,
+                    &String::from(js_sys::JsString::from(key)),
+                    py_val,
+                );
+            }
+        }
+
+        vm.ctx.set_item(vars, "js_vars", injections);
+    });
+
+    res.map(|value| py_to_js(&mut vm, value))
         .map_err(|err| py_str_err(&mut vm, &err).into())
 }
 
@@ -81,7 +130,7 @@ pub fn run_from_textbox(source: &str) -> Result<JsValue, JsValue> {
         vm.context().new_rustfunc(wasm_builtins::builtin_print_html),
     );
 
-    match eval(&mut vm, source) {
+    match eval(&mut vm, source, |_, _| {}) {
         Ok(value) => {
             console::log_1(&"Execution successful".into());
             match value.borrow().kind {
