diff options
| -rw-r--r-- | py/vm.c | 4 | ||||
| -rw-r--r-- | tests/basics/gen_yield_from.py | 13 |
2 files changed, 16 insertions, 1 deletions
@@ -1162,6 +1162,7 @@ yield: mp_obj_t send_value = POP(); mp_obj_t t_exc = MP_OBJ_NULL; mp_obj_t ret_value; + code_state->sp = sp; // Save sp because it's needed if mp_resume raises StopIteration if (inject_exc != MP_OBJ_NULL) { t_exc = inject_exc; inject_exc = MP_OBJ_NULL; @@ -1361,7 +1362,8 @@ exception_handler: } 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(MP_OBJ_FROM_PTR(nlr.ret_val)); + // (Instead of stack pop then push we just replace exhausted gen with value) + *code_state->sp = mp_obj_exception_get_value(MP_OBJ_FROM_PTR(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.py b/tests/basics/gen_yield_from.py index 5196b48d2..4e68aec63 100644 --- a/tests/basics/gen_yield_from.py +++ b/tests/basics/gen_yield_from.py @@ -40,3 +40,16 @@ def gen6(): g = gen6() print(list(g)) + +# StopIteration from within a Python function, within a native iterator (map), within a yield from +def gen7(x): + if x < 3: + return x + else: + raise StopIteration(444) + +def gen8(): + print((yield from map(gen7, range(100)))) + +g = gen8() +print(list(g)) |