Await jsproxy (pyodide#880)

Co-authored-by: Wei Ouyang <oeway007@gmail.com>
jmsmdy · Jan 10, 2021 · 7b45762 · 7b45762
1 parent 7255504
commit 7b45762
Show file tree

Hide file tree

Showing 5 changed files with 297 additions and 11 deletions.
diff --git a/src/core/hiwire.c b/src/core/hiwire.c
@@ -96,13 +96,20 @@ EM_JS(int, hiwire_init, (), {
   {
     // clang-format off
     if ((idval & 1) === 0) {
+      // clang-format on
       // least significant bit unset ==> idval is a singleton.
       // We don't reference count singletons.
-      // clang-format on
       return;
     }
     _hiwire.objects.delete(idval);
   };
+
+  Module.hiwire.isPromise = function(obj)
+  {
+    // clang-format off
+    return Object.prototype.toString.call(obj) === "[object Promise]";
+    // clang-format on
+  };
   return 0;
 });
 
@@ -369,6 +376,21 @@ EM_JS_NUM(bool, hiwire_is_function, (JsRef idobj), {
   // clang-format on
 });
 
+EM_JS_NUM(bool, hiwire_is_promise, (JsRef idobj), {
+  // clang-format off
+  let obj = Module.hiwire.get_value(idobj);
+  return Module.hiwire.isPromise(obj);
+  // clang-format on
+});
+
+EM_JS_REF(JsRef, hiwire_resolve_promise, (JsRef idobj), {
+  // clang-format off
+  let obj = Module.hiwire.get_value(idobj);
+  let result = Promise.resolve(obj);
+  return Module.hiwire.new_value(result);
+  // clang-format on
+});
+
 EM_JS_REF(JsRef, hiwire_to_string, (JsRef idobj), {
   return Module.hiwire.new_value(Module.hiwire.get_value(idobj).toString());
 });

diff --git a/src/core/hiwire.h b/src/core/hiwire.h
@@ -449,6 +449,20 @@ hiwire_get_bool(JsRef idobj);
 bool
 hiwire_is_function(JsRef idobj);
 
+/**
+ * Returns true if the object is a promise.
+ */
+bool
+hiwire_is_promise(JsRef idobj);
+
+/**
+ * Returns Promise.resolve(obj)
+ *
+ * Returns: New reference to Javascript promise
+ */
+JsRef
+hiwire_resolve_promise(JsRef idobj);
+
 /**
  * Gets the string representation of an object by calling `toString`.
  *

diff --git a/src/core/jsproxy.c b/src/core/jsproxy.c
@@ -1,14 +1,21 @@
 #define PY_SSIZE_T_CLEAN
 #include "Python.h"
 
-#include "jsproxy.h"
-
 #include "hiwire.h"
 #include "js2python.h"
+#include "jsproxy.h"
 #include "python2js.h"
 
 #include "structmember.h"
 
+_Py_IDENTIFIER(get_event_loop);
+_Py_IDENTIFIER(create_future);
+_Py_IDENTIFIER(set_exception);
+_Py_IDENTIFIER(set_result);
+_Py_IDENTIFIER(__await__);
+
+static PyObject* asyncio_get_event_loop;
+
 static PyTypeObject* PyExc_BaseException_Type;
 
 _Py_IDENTIFIER(__dir__);
@@ -19,7 +26,7 @@ JsBoundMethod_cnew(JsRef this_, const char* name);
 ////////////////////////////////////////////////////////////
 // JsProxy
 //
-// This is a Python object that provides ideomatic access to a Javascript
+// This is a Python object that provides idiomatic access to a Javascript
 // object.
 
 // clang-format off
@@ -28,6 +35,7 @@ typedef struct
   PyObject_HEAD
   JsRef js;
   PyObject* bytes;
+  bool awaited; // for promises
 } JsProxy;
 // clang-format on
 
@@ -37,7 +45,7 @@ static void
 JsProxy_dealloc(JsProxy* self)
 {
   hiwire_decref(self->js);
-  Py_XDECREF(self->bytes);
+  Py_CLEAR(self->bytes);
   Py_TYPE(self)->tp_free((PyObject*)self);
 }
 
@@ -447,6 +455,67 @@ JsProxy_Bool(PyObject* o)
   return hiwire_get_bool(self->js) ? 1 : 0;
 }
 
+PyObject*
+JsProxy_Await(JsProxy* self)
+{
+  // Guards
+  if (self->awaited) {
+    PyErr_SetString(PyExc_RuntimeError,
+                    "cannot reuse already awaited coroutine");
+    return NULL;
+  }
+
+  if (!hiwire_is_promise(self->js)) {
+    PyObject* str = JsProxy_Repr((PyObject*)self);
+    const char* str_utf8 = PyUnicode_AsUTF8(str);
+    PyErr_Format(PyExc_TypeError,
+                 "object %s can't be used in 'await' expression",
+                 str_utf8);
+    return NULL;
+  }
+
+  // Main
+  PyObject* result = NULL;
+
+  PyObject* loop = NULL;
+  PyObject* fut = NULL;
+  PyObject* set_result = NULL;
+  PyObject* set_exception = NULL;
+
+  loop = _PyObject_CallNoArg(asyncio_get_event_loop);
+  FAIL_IF_NULL(loop);
+
+  fut = _PyObject_CallMethodId(loop, &PyId_create_future, NULL);
+  FAIL_IF_NULL(fut);
+
+  set_result = _PyObject_GetAttrId(fut, &PyId_set_result);
+  FAIL_IF_NULL(set_result);
+  set_exception = _PyObject_GetAttrId(fut, &PyId_set_exception);
+  FAIL_IF_NULL(set_exception);
+
+  JsRef promise_id = hiwire_resolve_promise(self->js);
+  JsRef idargs = hiwire_array();
+  JsRef idarg;
+  // TODO: does this leak set_result and set_exception? See #1006.
+  idarg = python2js(set_result);
+  hiwire_push_array(idargs, idarg);
+  hiwire_decref(idarg);
+  idarg = python2js(set_exception);
+  hiwire_push_array(idargs, idarg);
+  hiwire_decref(idarg);
+  hiwire_decref(hiwire_call_member(promise_id, "then", idargs));
+  hiwire_decref(promise_id);
+  hiwire_decref(idargs);
+  result = _PyObject_CallMethodId(fut, &PyId___await__, NULL);
+
+finally:
+  Py_CLEAR(loop);
+  Py_CLEAR(set_result);
+  Py_CLEAR(set_exception);
+  Py_DECREF(fut);
+  return result;
+}
+
 // clang-format off
 static PyMappingMethods JsProxy_MappingMethods = {
   JsProxy_length,
@@ -472,6 +541,7 @@ static PyMethodDef JsProxy_Methods[] = {
     (PyCFunction)JsProxy_GetIter,
     METH_NOARGS,
     "Get an iterator over the object" },
+  { "__await__", (PyCFunction)JsProxy_Await, METH_NOARGS, ""},
   { "_has_bytes",
     (PyCFunction)JsProxy_HasBytes,
     METH_NOARGS,
@@ -484,6 +554,8 @@ static PyMethodDef JsProxy_Methods[] = {
 };
 // clang-format on
 
+static PyAsyncMethods JsProxy_asyncMethods = { .am_await =
+                                                 (unaryfunc)JsProxy_Await };
 static PyGetSetDef JsProxy_GetSet[] = { { "typeof", .get = JsProxy_typeof },
                                         { NULL } };
 
@@ -494,6 +566,7 @@ static PyTypeObject JsProxyType = {
   .tp_call = JsProxy_Call,
   .tp_getattro = JsProxy_GetAttr,
   .tp_setattro = JsProxy_SetAttr,
+  .tp_as_async = &JsProxy_asyncMethods,
   .tp_richcompare = JsProxy_RichCompare,
   .tp_flags = Py_TPFLAGS_DEFAULT,
   .tp_doc = "A proxy to make a Javascript object behave like a Python object",
@@ -514,6 +587,7 @@ JsProxy_cnew(JsRef idobj)
   self = (JsProxy*)JsProxyType.tp_alloc(&JsProxyType, 0);
   self->js = hiwire_incref(idobj);
   self->bytes = NULL;
+  self->awaited = false;
   return (PyObject*)self;
 }
 
@@ -702,23 +776,32 @@ int
 JsProxy_init()
 {
   bool success = false;
+
+  PyObject* asyncio_module = NULL;
+  PyObject* pyodide_module = NULL;
+
   PyExc_BaseException_Type = (PyTypeObject*)PyExc_BaseException;
   _Exc_JsException.tp_base = (PyTypeObject*)PyExc_Exception;
 
-  PyObject* module;
-  PyObject* exc;
+  asyncio_module = PyImport_ImportModule("asyncio");
+  FAIL_IF_NULL(asyncio_module);
+
+  asyncio_get_event_loop =
+    _PyObject_GetAttrId(asyncio_module, &PyId_get_event_loop);
+  FAIL_IF_NULL(asyncio_get_event_loop);
 
   // Add JsException to the pyodide module so people can catch it if they want.
-  module = PyImport_ImportModule("pyodide");
-  FAIL_IF_NULL(module);
+  pyodide_module = PyImport_ImportModule("pyodide");
+  FAIL_IF_NULL(pyodide_module);
   FAIL_IF_MINUS_ONE(
-    PyObject_SetAttrString(module, "JsException", Exc_JsException));
+    PyObject_SetAttrString(pyodide_module, "JsException", Exc_JsException));
   FAIL_IF_MINUS_ONE(PyType_Ready(&JsProxyType));
   FAIL_IF_MINUS_ONE(PyType_Ready(&JsBoundMethodType));
   FAIL_IF_MINUS_ONE(PyType_Ready(&_Exc_JsException));
 
   success = true;
 finally:
-  Py_CLEAR(module);
+  Py_CLEAR(asyncio_module);
+  Py_CLEAR(pyodide_module);
   return success ? 0 : -1;
 }
diff --git a/src/tests/test_jsproxy.py b/src/tests/test_jsproxy.py
@@ -1,4 +1,5 @@
 # See also test_typeconversions, and test_python.
+import pytest
 
 
 def test_jsproxy_dir(selenium):
@@ -189,3 +190,123 @@ def test_jsproxy_kwargs(selenium):
         )
         == 5
     )
+
+
+import time
+
+ASYNCIO_EVENT_LOOP_STARTUP = """
+import asyncio
+class DumbLoop(asyncio.AbstractEventLoop):
+    def create_future(self):
+        fut = asyncio.Future(loop=self)
+        old_set_result = fut.set_result
+        old_set_exception = fut.set_exception
+        def set_result(a):
+            print("set_result:", a)
+            old_set_result(a)
+        fut.set_result = set_result
+        def set_exception(a):
+            print("set_exception:", a)
+            old_set_exception(a)
+        fut.set_exception = set_exception
+        return fut
+
+    def get_debug(self):
+        return False
+
+asyncio.set_event_loop(DumbLoop())
+"""
+
+
+def test_await_jsproxy(selenium):
+    selenium.run(ASYNCIO_EVENT_LOOP_STARTUP)
+    selenium.run(
+        """
+        def prom(res,rej):
+            global resolve
+            resolve = res
+        from js import Promise
+        p = Promise.new(prom)
+        async def temp():
+            x = await p
+            return x + 7
+        resolve(10)
+        c = temp()
+        r = c.send(None)
+        """
+    )
+    time.sleep(0.01)
+    msg = "StopIteration: 17"
+    with pytest.raises(selenium.JavascriptException, match=msg):
+        selenium.run(
+            """
+            c.send(r.result())
+            """
+        )
+
+
+def test_await_fetch(selenium):
+    selenium.run(ASYNCIO_EVENT_LOOP_STARTUP)
+    selenium.run(
+        """
+        from js import fetch, window
+        async def test():
+            response = await fetch("console.html")
+            result = await response.text()
+            print(result)
+            return result
+        fetch = fetch.bind(window)
+
+        c = test()
+        r1 = c.send(None)
+        """
+    )
+    time.sleep(0.1)
+    selenium.run(
+        """
+        r2 = c.send(r1.result())
+        """
+    )
+    time.sleep(0.1)
+    msg = "StopIteration: <!doctype html>"
+    with pytest.raises(selenium.JavascriptException, match=msg):
+        selenium.run(
+            """
+            c.send(r2.result())
+            """
+        )
+
+
+def test_await_error(selenium):
+    selenium.run_js(
+        """
+        async function async_js_raises(){
+            console.log("Hello there???");
+            throw new Error("This is an error message!");
+        }
+        window.async_js_raises = async_js_raises;
+        function js_raises(){
+            throw new Error("This is an error message!");
+        }
+        window.js_raises = js_raises;
+        """
+    )
+    selenium.run(ASYNCIO_EVENT_LOOP_STARTUP)
+    selenium.run(
+        """
+        from js import async_js_raises, js_raises
+        async def test():
+            c = await async_js_raises()
+            return c
+        c = test()
+        r1 = c.send(None)
+        """
+    )
+    msg = "This is an error message!"
+    with pytest.raises(selenium.JavascriptException, match=msg):
+        # Wait for event loop to go around for chome
+        selenium.run(
+            """
+            r2 = c.send(r1.result())
+            """
+        )