Implement iterator / iterable / generator for pyproxy of Python itera…

…tors / iterables / generators
pyodide · Jan 31, 2021 · 2c95d38 · 2c95d38
1 parent 4de39e3
commit 2c95d38
Show file tree

Hide file tree

Showing 3 changed files with 218 additions and 17 deletions.
diff --git a/src/core/hiwire.c b/src/core/hiwire.c
@@ -110,6 +110,13 @@ EM_JS(int, hiwire_init, (), {
     _hiwire.objects.delete(idval);
   };
 
+  Module.hiwire.pop_value = function(idval)
+  {
+    let result = Module.hiwire.get_value(idval);
+    Module.hiwire.decref(idval);
+    return result;
+  };
+
   Module.hiwire.isPromise = function(obj)
   {
     // clang-format off

diff --git a/src/core/pyproxy.c b/src/core/pyproxy.c
@@ -138,6 +138,84 @@ _pyproxy_apply(PyObject* pyobj, JsRef idargs)
   return idresult;
 }
 
+// Return 2 if obj is iterator
+// Return 1 if iterable but not iterator
+// Return 0 if not iterable
+int
+_pyproxy_iterator_type(PyObject* obj)
+{
+  if (PyIter_Check(obj)) {
+    return 2;
+  }
+  PyObject* iter = PyObject_GetIter(obj);
+  int result = iter != NULL;
+  Py_CLEAR(iter);
+  PyErr_Clear();
+  return result;
+}
+
+JsRef
+_pyproxy_iter_next(PyObject* iterator)
+{
+  PyObject* item = PyIter_Next(iterator);
+  if (item == NULL) {
+    return NULL;
+  }
+  JsRef result = python2js(item);
+  Py_CLEAR(item);
+  return result;
+}
+
+JsRef
+_pyproxy_iter_send(PyObject* receiver, JsRef jsval)
+{
+  bool success = false;
+  PyObject* v = NULL;
+  PyObject* retval = NULL;
+  JsRef jsresult = NULL;
+
+  // cf implementation of YIELD_FROM opcode in ceval.c
+  v = js2python(jsval);
+  FAIL_IF_NULL(v);
+  if (PyGen_CheckExact(receiver) || PyCoro_CheckExact(receiver)) {
+    retval = _PyGen_Send((PyGenObject*)receiver, v);
+  } else if (v == Py_None) {
+    retval = Py_TYPE(receiver)->tp_iternext(receiver);
+  } else {
+    _Py_IDENTIFIER(send);
+    retval = _PyObject_CallMethodIdObjArgs(receiver, &PyId_send, v, NULL);
+  }
+  FAIL_IF_NULL(retval);
+
+  jsresult = python2js(retval);
+  FAIL_IF_NULL(jsresult);
+
+  success = true;
+finally:
+  Py_CLEAR(v);
+  Py_CLEAR(retval);
+  if (!success) {
+    hiwire_CLEAR(jsresult);
+  }
+  return jsresult;
+}
+
+JsRef
+_pyproxy_iter_fetch_stopiteration()
+{
+  PyObject* val = NULL;
+  // cf implementation of YIELD_FROM opcode in ceval.c
+  // _PyGen_FetchStopIterationValue returns an error code, but it seems
+  // redundant
+  _PyGen_FetchStopIterationValue(&val);
+  if (val == NULL) {
+    return NULL;
+  }
+  JsRef result = python2js(val);
+  Py_CLEAR(val);
+  return result;
+}
+
 void
 _pyproxy_destroy(PyObject* ptrobj)
 { // See bug #1049
@@ -161,6 +239,13 @@ EM_JS_REF(JsRef, pyproxy_new, (PyObject * ptrobj), {
   target['$$'] = { ptr : ptrobj, type : 'PyProxy' };
   Object.assign(target, Module.PyProxyPublicMethods);
   let proxy = new Proxy(target, Module.PyProxyHandlers);
+  let itertype = __pyproxy_iterator_type(ptrobj);
+  if (itertype == = 2) {
+    Object.assign(target, Module.PyProxyIteratorMethods);
+  }
+  if (itertype == = 1) {
+    Object.assign(target, Module.PyProxyIterableMethods);
+  }
   Module.PyProxies[ptrobj] = proxy;
 
   return Module.hiwire.new_value(proxy);
@@ -196,9 +281,7 @@ EM_JS(int, pyproxy_init, (), {
       if(jsref_repr === 0){
         _pythonexc2js();
       }
-      let repr = Module.hiwire.get_value(jsref_repr);
-      Module.hiwire.decref(jsref_repr);
-      return repr;
+      return Module.hiwire.pop_value(jsref_repr);
     },
     destroy : function() {
       let ptrobj = _getPtr(this);
@@ -219,9 +302,59 @@ EM_JS(int, pyproxy_init, (), {
       if(idresult === 0){
         _pythonexc2js();
       }
-      let jsresult = Module.hiwire.get_value(idresult);
-      Module.hiwire.decref(idresult);
-      return jsresult;
+      return Module.hiwire.pop_value(idresult);
+    },
+  };
+
+  // See: 
+  // https://docs.python.org/3/c-api/iter.html
+  // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Iteration_protocols
+  // This avoids allocating a PyProxy wrapper for the temporary iterator.
+  Module.PyProxyIterableMethods = {
+    [Symbol.iterator] : function*() {
+      let iterptr = _PyObject_GetIter(_getPtr(this));
+      if(iterptr === 0){
+        pythonexc2js();
+      }
+      let item;
+      while((item = __pyproxy_iter_next(iterptr))){
+        yield Module.hiwire.pop_value(item);
+      }
+      if(_PyErr_Occurred()){
+        pythonexc2js();
+      }
+      _Py_DecRef(iterptr);
+    }
+  };
+
+  Module.PyProxyIteratorMethods = {
+    [Symbol.iterator] : function() {
+      return this;
+    },    
+    next : function(arg) {
+      let idresult;
+      // Note: arg is optional, if arg is not supplied, it will be undefined
+      // which gets converted to "Py_None". This is as intended.
+      let idarg = Module.hiwire.new_value(arg);
+      try {
+        idresult = __pyproxy_iter_send(_getPtr(this), idarg);
+      } catch(e) {
+        Module.fatal_error(e);
+      } finally {
+        Module.hiwire.decref(idarg);
+      }      
+
+      let done = false;
+      if(idresult === 0){
+        idresult = __pyproxy_iter_fetch_stopiteration();
+        if (idresult){
+          done = true;
+        } else {
+          _pythonexc2js();
+        }
+      }
+      let value = Module.hiwire.pop_value(idresult);
+      return { done, value };
     },
   };
 
@@ -267,9 +400,7 @@ EM_JS(int, pyproxy_init, (), {
       if(idresult === 0){
         _pythonexc2js();
       }
-      let jsresult = Module.hiwire.get_value(idresult);
-      Module.hiwire.decref(idresult);
-      return jsresult;
+      return Module.hiwire.pop_value(idresult);
     },
     set: function (jsobj, jskey, jsval) {
       if(Reflect.has(jsobj, jskey) && !ignoredTargetFields.includes(jskey)){
@@ -290,9 +421,7 @@ EM_JS(int, pyproxy_init, (), {
       if(idresult === 0){
         _pythonexc2js();
       }
-      let jsresult = Module.hiwire.get_value(idresult);
-      Module.hiwire.decref(idresult);
-      return jsresult;
+      return Module.hiwire.pop_value(idresult);
     },
     deleteProperty: function (jsobj, jskey) {
       if(Reflect.has(jsobj, jskey) && !ignoredTargetFields.includes(jskey)){
@@ -311,9 +440,7 @@ EM_JS(int, pyproxy_init, (), {
       if(idresult === 0){
         _pythonexc2js();
       }
-      let jsresult = Module.hiwire.get_value(idresult);
-      Module.hiwire.decref(idresult);
-      return jsresult;
+      return Module.hiwire.pop_value(idresult);
     },
     ownKeys: function (jsobj) {
       let result = new Set(Reflect.ownKeys(jsobj));
@@ -327,8 +454,7 @@ EM_JS(int, pyproxy_init, (), {
       } catch(e) {
         Module.fatal_error(e);
       }
-      let jsresult = Module.hiwire.get_value(idresult);
-      Module.hiwire.decref(idresult);
+      let jsresult = Module.hiwire.pop_value(idresult);
       for(let key of jsresult){
         result.add(key);
       }

diff --git a/src/tests/test_pyproxy.py b/src/tests/test_pyproxy.py
@@ -132,3 +132,71 @@ def get_value(self, value):
             f.get_value();
             """
         )
+
+
+def test_pyproxy_iter(selenium):
+    assert (
+        selenium.run_js(
+            """
+        let c = pyodide.runPython(`
+            def test():
+                for i in range(10):
+                    yield i
+            test()
+        `);
+        return [...c];
+        """
+        )
+        == list(range(10))
+    )
+
+    assert (
+        set(
+            selenium.run_js(
+                """
+        c = pyodide.runPython(`
+            from collections import ChainMap
+            ChainMap({"a" : 2, "b" : 3})
+        `);
+        return [...c];
+        """
+            )
+        )
+        == set(["a", "b"])
+    )
+
+    [result, result2] = selenium.run_js(
+        """
+        let c = pyodide.runPython(`
+            def test():
+                acc = 0
+                for i in range(10):
+                    r = yield acc
+                    acc += i * r
+            test()
+        `)
+        let {done, value} = c.next();
+        let result = [];
+        while(!done){
+            result.push(value);
+            ({done, value} = c.next(value + 1));
+        }
+
+        function* test(){
+            let acc = 0;
+            for(let i=0; i < 10; i++){
+                let r = yield acc;
+                acc += i * r;
+            }
+        }
+        c = test();
+        ({done, value} = c.next());
+        let result2 = [];
+        while(!done){
+            result2.push(value);
+            ({done, value} = c.next(value + 1));
+        }
+        return [result, result2];
+        """
+    )
+    assert result == result2