webassembly/objpyproxy: Implement JS iterator protocol for Py iterables.

This allows using JavaScript for..of on Python iterables.

Signed-off-by: Damien George <damien@micropython.org>

5 files changed, 97 insertions, 4 deletions

diff --git a/ports/webassembly/Makefile b/ports/webassembly/Makefile
index 93b92ef58..fddfb5094 100644
--- a/ports/webassembly/Makefile
+++ b/ports/webassembly/Makefile
@@ -60,8 +60,10 @@ EXPORTED_FUNCTIONS_EXTRA += ,\
 	_proxy_c_to_js_dir,\
 	_proxy_c_to_js_get_array,\
 	_proxy_c_to_js_get_dict,\
+	_proxy_c_to_js_get_iter,\
 	_proxy_c_to_js_get_type,\
 	_proxy_c_to_js_has_attr,\
+	_proxy_c_to_js_iternext,\
 	_proxy_c_to_js_lookup_attr,\
 	_proxy_c_to_js_resume,\
 	_proxy_c_to_js_store_attr,\
diff --git a/ports/webassembly/objpyproxy.js b/ports/webassembly/objpyproxy.js
index 9ba06283e..3b94f8aad 100644
--- a/ports/webassembly/objpyproxy.js
+++ b/ports/webassembly/objpyproxy.js
@@ -162,10 +162,37 @@ const py_proxy_handler = {
         if (prop === "then") {
             return null;
         }
+
+        if (prop === Symbol.iterator) {
+            // Get the Python object iterator, and return a JavaScript generator.
+            const iter_ref = Module.ccall(
+                "proxy_c_to_js_get_iter",
+                "number",
+                ["number"],
+                [target._ref],
+            );
+            return function* () {
+                const value = Module._malloc(3 * 4);
+                while (true) {
+                    const valid = Module.ccall(
+                        "proxy_c_to_js_iternext",
+                        "number",
+                        ["number", "pointer"],
+                        [iter_ref, value],
+                    );
+                    if (!valid) {
+                        break;
+                    }
+                    yield proxy_convert_mp_to_js_obj_jsside(value);
+                }
+                Module._free(value);
+            };
+        }
+
         const value = Module._malloc(3 * 4);
         Module.ccall(
             "proxy_c_to_js_lookup_attr",
-            "number",
+            "null",
             ["number", "string", "pointer"],
             [target._ref, prop, value],
         );
diff --git a/ports/webassembly/proxy_c.c b/ports/webassembly/proxy_c.c
index ade9d36de..b874b36c0 100644
--- a/ports/webassembly/proxy_c.c
+++ b/ports/webassembly/proxy_c.c
@@ -65,6 +65,12 @@ void proxy_c_init(void) {
 
 MP_REGISTER_ROOT_POINTER(mp_obj_t proxy_c_ref);
 
+static inline size_t proxy_c_add_obj(mp_obj_t obj) {
+    size_t id = ((mp_obj_list_t *)MP_OBJ_TO_PTR(MP_STATE_PORT(proxy_c_ref)))->len;
+    mp_obj_list_append(MP_STATE_PORT(proxy_c_ref), obj);
+    return id;
+}
+
 static inline mp_obj_t proxy_c_get_obj(uint32_t c_ref) {
     return ((mp_obj_list_t *)MP_OBJ_TO_PTR(MP_STATE_PORT(proxy_c_ref)))->items[c_ref];
 }
@@ -122,9 +128,7 @@ void proxy_convert_mp_to_js_obj_cside(mp_obj_t obj, uint32_t *out) {
         } else {
             kind = PROXY_KIND_MP_OBJECT;
         }
-        size_t id = ((mp_obj_list_t *)MP_OBJ_TO_PTR(MP_STATE_PORT(proxy_c_ref)))->len;
-        mp_obj_list_append(MP_STATE_PORT(proxy_c_ref), obj);
-        out[1] = id;
+        out[1] = proxy_c_add_obj(obj);
     }
     out[0] = kind;
 }
@@ -285,6 +289,38 @@ void proxy_c_to_js_get_dict(uint32_t c_ref, uint32_t *out) {
 }
 
 /******************************************************************************/
+// Bridge Python iterator to JavaScript iterator protocol.
+
+uint32_t proxy_c_to_js_get_iter(uint32_t c_ref) {
+    mp_obj_t obj = proxy_c_get_obj(c_ref);
+    mp_obj_t iter = mp_getiter(obj, NULL);
+    return proxy_c_add_obj(iter);
+}
+
+bool proxy_c_to_js_iternext(uint32_t c_ref, uint32_t *out) {
+    mp_obj_t obj = proxy_c_get_obj(c_ref);
+    nlr_buf_t nlr;
+    if (nlr_push(&nlr) == 0) {
+        mp_obj_t iter = mp_iternext_allow_raise(obj);
+        if (iter == MP_OBJ_STOP_ITERATION) {
+            nlr_pop();
+            return false;
+        }
+        nlr_pop();
+        proxy_convert_mp_to_js_obj_cside(iter, out);
+        return true;
+    } else {
+        if (mp_obj_is_subclass_fast(MP_OBJ_FROM_PTR(((mp_obj_base_t *)nlr.ret_val)->type), MP_OBJ_FROM_PTR(&mp_type_StopIteration))) {
+            return false;
+        } else {
+            // uncaught exception
+            proxy_convert_mp_to_js_exc_cside(nlr.ret_val, out);
+            return true;
+        }
+    }
+}
+
+/******************************************************************************/
 // Bridge Python generator to JavaScript thenable.
 
 static const mp_obj_fun_builtin_var_t resume_obj;
diff --git a/tests/ports/webassembly/iterator.mjs b/tests/ports/webassembly/iterator.mjs
new file mode 100644
index 000000000..91415d530
--- /dev/null
+++ b/tests/ports/webassembly/iterator.mjs
@@ -0,0 +1,21 @@
+// Test accessing Python iterables from JavaScript via the JavaScript iterator protocol.
+
+const mp = await (await import(process.argv[2])).loadMicroPython();
+
+mp.runPython(`
+    s = "abc"
+    l = [1, 2, 3]
+`);
+
+// Iterate a Python string.
+for (const value of mp.globals.get("s")) {
+    console.log(value);
+}
+
+// Iterate a Python list.
+for (const value of mp.globals.get("l")) {
+    console.log(value);
+}
+
+// Iterate a Python list from a built-in JavaScript constructor.
+mp.runPython("import js; print(js.Set.new([1, 2, 3]).has(3))");
diff --git a/tests/ports/webassembly/iterator.mjs.exp b/tests/ports/webassembly/iterator.mjs.exp
new file mode 100644
index 000000000..9d1aaac11
--- /dev/null
+++ b/tests/ports/webassembly/iterator.mjs.exp
@@ -0,0 +1,7 @@
+a
+b
+c
+1
+2
+3
+True