vm: Properly handle StopIteration raised in user instance iterator.

I.e. in bytecode Python functions.

2 files changed, 28 insertions, 0 deletions

diff --git a/py/vm.c b/py/vm.c
index dc228bf0ff..aa84bdb0aa 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -1249,6 +1249,12 @@ exception_handler:
                         code_state->ip = ip + ulab; // jump to after for-block
                         code_state->sp -= 1; // pop the exhausted iterator
                         goto outer_dispatch_loop; // continue with dispatch loop
+                    } else if (*code_state->ip == MP_BC_YIELD_FROM) {
+                        // StopIteration inside yield from call means return a value of
+                        // yield from, so inject exception's value as yield from's result
+                        *++code_state->sp = mp_obj_exception_get_value(nlr.ret_val);
+                        code_state->ip++; // yield from is over, move to next instruction
+                        goto outer_dispatch_loop; // continue with dispatch loop
                     }
                 }
             }
diff --git a/tests/basics/gen_yield_from_ducktype.py b/tests/basics/gen_yield_from_ducktype.py
index aa0109c914..c02ec67562 100644
--- a/tests/basics/gen_yield_from_ducktype.py
+++ b/tests/basics/gen_yield_from_ducktype.py
@@ -42,3 +42,25 @@ g = gen3()
 print(next(g))
 print(g.send(5))
 print(g.send(100))
+
+
+#
+# Test proper handling of StopIteration vs other exceptions
+#
+class MyIter:
+    def __iter__(self):
+        return self
+    def __next__(self):
+        raise StopIteration(42)
+
+def gen4():
+    global ret
+    ret = yield from MyIter()
+    1//0
+
+ret = None
+try:
+    print(list(gen4()))
+except ZeroDivisionError:
+    print("ZeroDivisionError")
+print(ret)