Add PyProxy iterator support

Author: hoodmane

Lib/test/test_pyodide.py

@@ -390,7 +390,11 @@ class T:
             a = 7
             b = "zz"
 
-        l = set(x for x in run_js("(o) => Reflect.ownKeys(o)")(T) if isinstance(x, str))
+        l = set(
+            x
+            for x in run_js("(o) => Reflect.ownKeys(o)")(T)
+            if isinstance(x, str)
+        )
         self.assertGreaterEqual(l, {"a", "b"})
 
     def test_pyproxy_call_simple(self):
@@ -493,3 +497,17 @@ def test_pyproxy_set_del(self):
         self.assertEqual(d["y"], 3)
         run_js("(d) => d.delete('x')")(d)
         self.assertNotIn("x", d)
+
+    def test_pyproxy_iterator(self):
+        d = [7, 21, 39]
+        res = list(
+            run_js("(d) => [d.next(), d.next(), d.next(), d.next()]")(iter(d))
+        )
+        self.assertFalse(res[0].done)
+        self.assertFalse(res[1].done)
+        self.assertFalse(res[2].done)
+        self.assertTrue(res[3].done)
+        self.assertEqual(res[0].value, 7)
+        self.assertEqual(res[1].value, 21)
+        self.assertEqual(res[2].value, 39)
+        self.assertEqual(res[3].value, None)

Modules/_pyodide_core/pyproxy.c

@@ -12,6 +12,8 @@
 #define HAS_LENGTH             (1 << 2)
 #define HAS_SET                (1 << 3)
 #define IS_CALLABLE            (1 << 4)
+#define IS_ITERABLE            (1 << 7)
+#define IS_ITERATOR            (1 << 8)
 
 EM_JS_VAL(JsVal, pyproxy_new, (PyObject * ptrobj), {
   return Module.pyproxy_new(ptrobj);
@@ -80,6 +82,15 @@ type_getflags(PyTypeObject* obj_type)
   SET_FLAG_IF(HAS_LENGTH, seq_proto->sq_length || map_proto->mp_length);
   SET_FLAG_IF(HAS_SET, map_proto->mp_ass_subscript || seq_proto->sq_ass_item);
   SET_FLAG_IF(IS_CALLABLE, obj_type->tp_call);
+  SET_FLAG_IF(IS_ITERABLE, obj_type->tp_iter || seq_proto->sq_item);
+
+  extern PyObject* _PyObject_NextNotImplemented(PyObject *);
+  if (obj_type->tp_iternext != NULL &&
+      obj_type->tp_iternext != &_PyObject_NextNotImplemented) {
+    result &= ~IS_ITERABLE;
+    result |= IS_ITERATOR;
+  }
+
   return result;
 
 #undef SET_FLAG_IF
@@ -431,3 +442,32 @@ _pyproxy_apply(PyObject* callable,
   return result;
 }
 
+EM_JS(JsVal, _pyproxyGen_make_result, (bool done, JsVal value), {
+  return { done : !!done, value };
+})
+
+EMSCRIPTEN_KEEPALIVE JsVal
+_pyproxyGen_Send(PyObject* receiver, JsVal jsval)
+{
+  bool success = false;
+  PyObject* v = NULL;
+  PyObject* retval = NULL;
+
+  v = js2python(jsval);
+  FAIL_IF_NULL(v);
+  PySendResult status = PyIter_Send(receiver, v, &retval);
+  if (status == PYGEN_ERROR) {
+    FAIL();
+  }
+  JsVal result = python2js(retval);
+  FAIL_IF_JS_ERROR(result);
+
+  success = true;
+finally:
+  Py_CLEAR(v);
+  Py_CLEAR(retval);
+  if (!success) {
+    return JS_ERROR;
+  }
+  return _pyproxyGen_make_result(status == PYGEN_RETURN, result);
+}

Modules/_pyodide_core/pyproxy.ts

@@ -34,6 +34,7 @@ declare function __pyproxy_apply(
   kwargs_names: string[],
   num_kwargs: number,
 ): any;
+declare function __pyproxyGen_Send(ptr: number, arg: any): IteratorResult<any>
 
 // pyodide-skip
 
@@ -298,6 +299,7 @@ function getPyProxyClass(flags: number) {
     [HAS_LENGTH, PyLengthMethods],
     [HAS_SET, PySetItemMethods],
     [IS_CALLABLE, PyCallableMethods],
+    [IS_ITERATOR, PyIteratorMethods],
   ];
   for (let [feature_flag, methods] of FLAG_TYPE_PAIRS) {
     if (flags & feature_flag) {
@@ -1038,3 +1040,57 @@ function callPyObjectKwargs(ptrobj: number, jsargs: any[], kwargs: any) {
 function callPyObject(ptrobj: number, jsargs: any) {
   return callPyObjectKwargs(ptrobj, jsargs, {});
 }
+
+
+/**
+ * A :js:class:`~pyodide.ffi.PyProxy` whose proxied Python object is an :term:`iterator`
+ * (i.e., has a :meth:`~generator.send` or :meth:`~iterator.__next__` method).
+ */
+class PyIterator extends PyProxy {
+  /** @private */
+  static [Symbol.hasInstance](obj: any): obj is PyProxy {
+    return API.isPyProxy(obj) && !!(_getFlags(obj) & IS_ITERATOR);
+  }
+}
+
+interface PyIterator extends PyIteratorMethods {}
+
+// Controlled by IS_ITERATOR, appears for any object with a __next__ or
+// tp_iternext method.
+class PyIteratorMethods {
+  /** @private */
+  [Symbol.iterator]() {
+    return this;
+  }
+  /**
+   * This translates to the Python code ``next(obj)``. Returns the next value of
+   * the generator. See the documentation for :js:meth:`Generator.next` The
+   * argument will be sent to the Python generator.
+   *
+   * This will be used implicitly by ``for(let x of proxy){}``.
+   *
+   * @param arg The value to send to the generator. The value will be assigned
+   * as a result of a yield expression.
+   * @returns An Object with two properties: ``done`` and ``value``. When the
+   * generator yields ``some_value``, ``next`` returns ``{done : false, value :
+   * some_value}``. When the generator raises a :py:exc:`StopIteration`
+   * exception, ``next`` returns ``{done : true, value : result_value}``.
+   */
+  next(arg: any = undefined): IteratorResult<any, any> {
+    // Note: arg is optional, if arg is not supplied, it will be undefined
+    // which gets converted to "Py_None". This is as intended.
+    let result;
+    let done;
+    try {
+      Py_ENTER();
+      result = __pyproxyGen_Send(_getPtr(this), arg);
+      Py_EXIT();
+    } catch (e) {
+      API.fatal_error(e);
+    }
+    if (result === Module.error) {
+      _pythonexc2js();
+    }
+    return result;
+  }
+}