diff options
Diffstat (limited to 'py/objgenerator.c')
| -rw-r--r-- | py/objgenerator.c | 30 |
1 files changed, 12 insertions, 18 deletions
diff --git a/py/objgenerator.c b/py/objgenerator.c index 1ee7b8b1d..784310092 100644 --- a/py/objgenerator.c +++ b/py/objgenerator.c @@ -238,16 +238,20 @@ mp_vm_return_kind_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_ return ret_kind; } -STATIC mp_obj_t gen_resume_and_raise(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value) { +STATIC mp_obj_t gen_resume_and_raise(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value, bool raise_stop_iteration) { mp_obj_t ret; switch (mp_obj_gen_resume(self_in, send_value, throw_value, &ret)) { case MP_VM_RETURN_NORMAL: default: - // Optimize return w/o value in case generator is used in for loop - if (ret == mp_const_none || ret == MP_OBJ_STOP_ITERATION) { - return MP_OBJ_STOP_ITERATION; + // A normal return is a StopIteration, either raise it or return + // MP_OBJ_STOP_ITERATION as an optimisation. + if (ret == mp_const_none) { + ret = MP_OBJ_NULL; + } + if (raise_stop_iteration) { + mp_raise_StopIteration(ret); } else { - nlr_raise(mp_obj_new_exception_arg1(&mp_type_StopIteration, ret)); + return mp_make_stop_iteration(ret); } case MP_VM_RETURN_YIELD: @@ -259,16 +263,11 @@ STATIC mp_obj_t gen_resume_and_raise(mp_obj_t self_in, mp_obj_t send_value, mp_o } STATIC mp_obj_t gen_instance_iternext(mp_obj_t self_in) { - return gen_resume_and_raise(self_in, mp_const_none, MP_OBJ_NULL); + return gen_resume_and_raise(self_in, mp_const_none, MP_OBJ_NULL, false); } STATIC mp_obj_t gen_instance_send(mp_obj_t self_in, mp_obj_t send_value) { - mp_obj_t ret = gen_resume_and_raise(self_in, send_value, MP_OBJ_NULL); - if (ret == MP_OBJ_STOP_ITERATION) { - mp_raise_type(&mp_type_StopIteration); - } else { - return ret; - } + return gen_resume_and_raise(self_in, send_value, MP_OBJ_NULL, true); } STATIC MP_DEFINE_CONST_FUN_OBJ_2(gen_instance_send_obj, gen_instance_send); @@ -290,12 +289,7 @@ STATIC mp_obj_t gen_instance_throw(size_t n_args, const mp_obj_t *args) { exc = args[2]; } - mp_obj_t ret = gen_resume_and_raise(args[0], mp_const_none, exc); - if (ret == MP_OBJ_STOP_ITERATION) { - mp_raise_type(&mp_type_StopIteration); - } else { - return ret; - } + return gen_resume_and_raise(args[0], mp_const_none, exc, true); } STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(gen_instance_throw_obj, 2, 4, gen_instance_throw); |