diff options
Diffstat (limited to 'py/runtime.c')
| -rw-r--r-- | py/runtime.c | 20 |
1 files changed, 17 insertions, 3 deletions
diff --git a/py/runtime.c b/py/runtime.c index b53711bbe..2c849fe95 100644 --- a/py/runtime.c +++ b/py/runtime.c @@ -1217,6 +1217,7 @@ mp_obj_t mp_getiter(mp_obj_t o_in, mp_obj_iter_buf_t *iter_buf) { mp_obj_t mp_iternext_allow_raise(mp_obj_t o_in) { const mp_obj_type_t *type = mp_obj_get_type(o_in); if (type->iternext != NULL) { + MP_STATE_THREAD(stop_iteration_arg) = MP_OBJ_NULL; return type->iternext(o_in); } else { // check for __next__ method @@ -1242,6 +1243,7 @@ mp_obj_t mp_iternext(mp_obj_t o_in) { MP_STACK_CHECK(); // enumerate, filter, map and zip can recursively call mp_iternext const mp_obj_type_t *type = mp_obj_get_type(o_in); if (type->iternext != NULL) { + MP_STATE_THREAD(stop_iteration_arg) = MP_OBJ_NULL; return type->iternext(o_in); } else { // check for __next__ method @@ -1256,7 +1258,7 @@ mp_obj_t mp_iternext(mp_obj_t o_in) { return ret; } 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 MP_OBJ_STOP_ITERATION; + return mp_make_stop_iteration(mp_obj_exception_get_value(MP_OBJ_FROM_PTR(nlr.ret_val))); } else { nlr_jump(nlr.ret_val); } @@ -1281,14 +1283,18 @@ mp_vm_return_kind_t mp_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t th } if (type->iternext != NULL && send_value == mp_const_none) { + MP_STATE_THREAD(stop_iteration_arg) = MP_OBJ_NULL; mp_obj_t ret = type->iternext(self_in); *ret_val = ret; if (ret != MP_OBJ_STOP_ITERATION) { return MP_VM_RETURN_YIELD; } else { // The generator is finished. - // This is an optimised "raise StopIteration(None)". - *ret_val = mp_const_none; + // This is an optimised "raise StopIteration(*ret_val)". + *ret_val = MP_STATE_THREAD(stop_iteration_arg); + if (*ret_val == MP_OBJ_NULL) { + *ret_val = mp_const_none; + } return MP_VM_RETURN_NORMAL; } } @@ -1559,6 +1565,14 @@ NORETURN void mp_raise_NotImplementedError(mp_rom_error_text_t msg) { #endif +NORETURN void mp_raise_StopIteration(mp_obj_t arg) { + if (arg == MP_OBJ_NULL) { + mp_raise_type(&mp_type_StopIteration); + } else { + nlr_raise(mp_obj_new_exception_arg1(&mp_type_StopIteration, arg)); + } +} + NORETURN void mp_raise_OSError(int errno_) { nlr_raise(mp_obj_new_exception_arg1(&mp_type_OSError, MP_OBJ_NEW_SMALL_INT(errno_))); } |