py: Support single argument to optimised MP_OBJ_STOP_ITERATION.

The MP_OBJ_STOP_ITERATION optimisation is a shortcut for creating a
StopIteration() exception object, and means that heap memory does not need
to be allocated for the exception (in cases where it can be used).  This
commit allows this optimised object to take an optional argument (before,
it could only have no argument).

The commit also adds some new tests to cover corner cases with
StopIteration and generators that previously did not work.

Signed-off-by: Damien George <damien@micropython.org>
This commit is contained in:
Damien George 2021-06-29 17:34:34 +10:00
parent e3825e28e6
commit bb00125aaa
8 changed files with 115 additions and 24 deletions

View File

@ -322,7 +322,7 @@ STATIC mp_obj_t mp_builtin_next(size_t n_args, const mp_obj_t *args) {
if (n_args == 1) { if (n_args == 1) {
mp_obj_t ret = mp_iternext_allow_raise(args[0]); mp_obj_t ret = mp_iternext_allow_raise(args[0]);
if (ret == MP_OBJ_STOP_ITERATION) { if (ret == MP_OBJ_STOP_ITERATION) {
mp_raise_type(&mp_type_StopIteration); mp_raise_StopIteration(MP_STATE_THREAD(stop_iteration_arg));
} else { } else {
return ret; return ret;
} }
@ -336,7 +336,7 @@ MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_builtin_next_obj, 1, 2, mp_builtin_next);
STATIC mp_obj_t mp_builtin_next(mp_obj_t o) { STATIC mp_obj_t mp_builtin_next(mp_obj_t o) {
mp_obj_t ret = mp_iternext_allow_raise(o); mp_obj_t ret = mp_iternext_allow_raise(o);
if (ret == MP_OBJ_STOP_ITERATION) { if (ret == MP_OBJ_STOP_ITERATION) {
mp_raise_type(&mp_type_StopIteration); mp_raise_StopIteration(MP_STATE_THREAD(stop_iteration_arg));
} else { } else {
return ret; return ret;
} }

View File

@ -266,6 +266,9 @@ typedef struct _mp_state_thread_t {
// pending exception object (MP_OBJ_NULL if not pending) // pending exception object (MP_OBJ_NULL if not pending)
volatile mp_obj_t mp_pending_exception; volatile mp_obj_t mp_pending_exception;
// If MP_OBJ_STOP_ITERATION is propagated then this holds its argument.
mp_obj_t stop_iteration_arg;
#if MICROPY_PY_SYS_SETTRACE #if MICROPY_PY_SYS_SETTRACE
mp_obj_t prof_trace_callback; mp_obj_t prof_trace_callback;
bool prof_callback_is_executing; bool prof_callback_is_executing;

View File

@ -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; 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; mp_obj_t ret;
switch (mp_obj_gen_resume(self_in, send_value, throw_value, &ret)) { switch (mp_obj_gen_resume(self_in, send_value, throw_value, &ret)) {
case MP_VM_RETURN_NORMAL: case MP_VM_RETURN_NORMAL:
default: default:
// Optimize return w/o value in case generator is used in for loop // A normal return is a StopIteration, either raise it or return
if (ret == mp_const_none || ret == MP_OBJ_STOP_ITERATION) { // MP_OBJ_STOP_ITERATION as an optimisation.
return MP_OBJ_STOP_ITERATION; if (ret == mp_const_none) {
ret = MP_OBJ_NULL;
}
if (raise_stop_iteration) {
mp_raise_StopIteration(ret);
} else { } else {
nlr_raise(mp_obj_new_exception_arg1(&mp_type_StopIteration, ret)); return mp_make_stop_iteration(ret);
} }
case MP_VM_RETURN_YIELD: 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) { 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) { 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); return gen_resume_and_raise(self_in, send_value, MP_OBJ_NULL, true);
if (ret == MP_OBJ_STOP_ITERATION) {
mp_raise_type(&mp_type_StopIteration);
} else {
return ret;
}
} }
STATIC MP_DEFINE_CONST_FUN_OBJ_2(gen_instance_send_obj, gen_instance_send); 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]; exc = args[2];
} }
mp_obj_t ret = gen_resume_and_raise(args[0], mp_const_none, exc); return gen_resume_and_raise(args[0], mp_const_none, exc, true);
if (ret == MP_OBJ_STOP_ITERATION) {
mp_raise_type(&mp_type_StopIteration);
} else {
return ret;
}
} }
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(gen_instance_throw_obj, 2, 4, gen_instance_throw); STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(gen_instance_throw_obj, 2, 4, gen_instance_throw);

View File

@ -48,7 +48,6 @@ STATIC mp_obj_t it_iternext(mp_obj_t self_in) {
// an exception was raised // an exception was raised
mp_obj_type_t *t = (mp_obj_type_t *)((mp_obj_base_t *)nlr.ret_val)->type; mp_obj_type_t *t = (mp_obj_type_t *)((mp_obj_base_t *)nlr.ret_val)->type;
if (t == &mp_type_StopIteration || t == &mp_type_IndexError) { if (t == &mp_type_StopIteration || t == &mp_type_IndexError) {
// return MP_OBJ_STOP_ITERATION instead of raising
return MP_OBJ_STOP_ITERATION; return MP_OBJ_STOP_ITERATION;
} else { } else {
// re-raise exception // re-raise exception

View File

@ -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) { mp_obj_t mp_iternext_allow_raise(mp_obj_t o_in) {
const mp_obj_type_t *type = mp_obj_get_type(o_in); const mp_obj_type_t *type = mp_obj_get_type(o_in);
if (type->iternext != NULL) { if (type->iternext != NULL) {
MP_STATE_THREAD(stop_iteration_arg) = MP_OBJ_NULL;
return type->iternext(o_in); return type->iternext(o_in);
} else { } else {
// check for __next__ method // 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 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); const mp_obj_type_t *type = mp_obj_get_type(o_in);
if (type->iternext != NULL) { if (type->iternext != NULL) {
MP_STATE_THREAD(stop_iteration_arg) = MP_OBJ_NULL;
return type->iternext(o_in); return type->iternext(o_in);
} else { } else {
// check for __next__ method // check for __next__ method
@ -1256,7 +1258,7 @@ mp_obj_t mp_iternext(mp_obj_t o_in) {
return ret; return ret;
} else { } 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))) { 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 { } else {
nlr_jump(nlr.ret_val); 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) { 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); mp_obj_t ret = type->iternext(self_in);
*ret_val = ret; *ret_val = ret;
if (ret != MP_OBJ_STOP_ITERATION) { if (ret != MP_OBJ_STOP_ITERATION) {
return MP_VM_RETURN_YIELD; return MP_VM_RETURN_YIELD;
} else { } else {
// The generator is finished. // The generator is finished.
// This is an optimised "raise StopIteration(None)". // This is an optimised "raise StopIteration(*ret_val)".
*ret_val = mp_const_none; *ret_val = MP_STATE_THREAD(stop_iteration_arg);
if (*ret_val == MP_OBJ_NULL) {
*ret_val = mp_const_none;
}
return MP_VM_RETURN_NORMAL; return MP_VM_RETURN_NORMAL;
} }
} }
@ -1559,6 +1565,14 @@ NORETURN void mp_raise_NotImplementedError(mp_rom_error_text_t msg) {
#endif #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_) { NORETURN void mp_raise_OSError(int errno_) {
nlr_raise(mp_obj_new_exception_arg1(&mp_type_OSError, MP_OBJ_NEW_SMALL_INT(errno_))); nlr_raise(mp_obj_new_exception_arg1(&mp_type_OSError, MP_OBJ_NEW_SMALL_INT(errno_)));
} }

View File

@ -154,6 +154,11 @@ mp_obj_t mp_iternext_allow_raise(mp_obj_t o); // may return MP_OBJ_STOP_ITERATIO
mp_obj_t mp_iternext(mp_obj_t o); // will always return MP_OBJ_STOP_ITERATION instead of raising StopIteration(...) mp_obj_t mp_iternext(mp_obj_t o); // will always return MP_OBJ_STOP_ITERATION instead of raising StopIteration(...)
mp_vm_return_kind_t mp_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value, mp_obj_t *ret_val); mp_vm_return_kind_t mp_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value, mp_obj_t *ret_val);
static inline mp_obj_t mp_make_stop_iteration(mp_obj_t o) {
MP_STATE_THREAD(stop_iteration_arg) = o;
return MP_OBJ_STOP_ITERATION;
}
mp_obj_t mp_make_raise_obj(mp_obj_t o); mp_obj_t mp_make_raise_obj(mp_obj_t o);
mp_obj_t mp_import_name(qstr name, mp_obj_t fromlist, mp_obj_t level); mp_obj_t mp_import_name(qstr name, mp_obj_t fromlist, mp_obj_t level);
@ -179,6 +184,7 @@ NORETURN void mp_raise_TypeError(mp_rom_error_text_t msg);
NORETURN void mp_raise_NotImplementedError(mp_rom_error_text_t msg); NORETURN void mp_raise_NotImplementedError(mp_rom_error_text_t msg);
#endif #endif
NORETURN void mp_raise_StopIteration(mp_obj_t arg);
NORETURN void mp_raise_OSError(int errno_); NORETURN void mp_raise_OSError(int errno_);
NORETURN void mp_raise_recursion_depth(void); NORETURN void mp_raise_recursion_depth(void);

View File

@ -16,3 +16,15 @@ try:
next(run()) next(run())
except StopIteration: except StopIteration:
print("StopIteration") print("StopIteration")
# Where "f" is a native generator
def run():
print((yield from f))
f = zip()
try:
next(run())
except StopIteration:
print("StopIteration")

View File

@ -0,0 +1,63 @@
# test StopIteration interaction with generators
try:
enumerate, exec
except:
print("SKIP")
raise SystemExit
def get_stop_iter_arg(msg, code):
try:
exec(code)
print("FAIL")
except StopIteration as er:
print(msg, er.args)
class A:
def __iter__(self):
return self
def __next__(self):
raise StopIteration(42)
class B:
def __getitem__(self, index):
# argument to StopIteration should get ignored
raise StopIteration(42)
def gen(x):
return x
yield
def gen2(x):
try:
yield
except ValueError:
pass
return x
get_stop_iter_arg("next", "next(A())")
get_stop_iter_arg("iter", "next(iter(B()))")
get_stop_iter_arg("enumerate", "next(enumerate(A()))")
get_stop_iter_arg("map", "next(map(lambda x:x, A()))")
get_stop_iter_arg("zip", "next(zip(A()))")
g = gen(None)
get_stop_iter_arg("generator0", "next(g)")
get_stop_iter_arg("generator1", "next(g)")
g = gen(42)
get_stop_iter_arg("generator0", "next(g)")
get_stop_iter_arg("generator1", "next(g)")
get_stop_iter_arg("send", "gen(None).send(None)")
get_stop_iter_arg("send", "gen(42).send(None)")
g = gen2(None)
next(g)
get_stop_iter_arg("throw", "g.throw(ValueError)")
g = gen2(42)
next(g)
get_stop_iter_arg("throw", "g.throw(ValueError)")