From 962b1cd1b120d777636ce8195d14f3d686e96619 Mon Sep 17 00:00:00 2001 From: Paul Sokolovsky Date: Sun, 23 Mar 2014 21:48:29 +0200 Subject: [PATCH] objgenerator: Implement return with value and .close() method. Return with value gets converted to StopIteration(value). Implementation keeps optimizing against creating of possibly unneeded exception objects, so there're considerable refactoring to implement these features. --- py/obj.h | 2 + py/objexcept.c | 14 +++++++ py/objgenerator.c | 71 ++++++++++++++++++++++++++------ py/objgenerator.h | 1 + tests/basics/generator-return.py | 10 +++++ tests/basics/generator_close.py | 59 ++++++++++++++++++++++++++ 6 files changed, 144 insertions(+), 13 deletions(-) create mode 100644 py/objgenerator.h create mode 100644 tests/basics/generator-return.py create mode 100644 tests/basics/generator_close.py diff --git a/py/obj.h b/py/obj.h index f80bd4352..9599d137d 100644 --- a/py/obj.h +++ b/py/obj.h @@ -273,6 +273,7 @@ mp_obj_t mp_obj_new_float(mp_float_t val); mp_obj_t mp_obj_new_complex(mp_float_t real, mp_float_t imag); #endif mp_obj_t mp_obj_new_exception(const mp_obj_type_t *exc_type); +mp_obj_t mp_obj_new_exception_args(const mp_obj_type_t *exc_type, uint n_args, const mp_obj_t *args); mp_obj_t mp_obj_new_exception_msg(const mp_obj_type_t *exc_type, const char *msg); mp_obj_t mp_obj_new_exception_msg_varg(const mp_obj_type_t *exc_type, const char *fmt, ...); // counts args by number of % symbols in fmt, excluding %%; can only handle void* sizes (ie no float/double!) mp_obj_t mp_obj_new_range(int start, int stop, int step); @@ -342,6 +343,7 @@ machine_int_t mp_obj_int_get_checked(mp_obj_t self_in); // exception bool mp_obj_is_exception_type(mp_obj_t self_in); bool mp_obj_is_exception_instance(mp_obj_t self_in); +bool mp_obj_exception_match(mp_obj_t exc, const mp_obj_type_t *exc_type); void mp_obj_exception_clear_traceback(mp_obj_t self_in); void mp_obj_exception_add_traceback(mp_obj_t self_in, qstr file, machine_uint_t line, qstr block); void mp_obj_exception_get_traceback(mp_obj_t self_in, machine_uint_t *n, machine_uint_t **values); diff --git a/py/objexcept.c b/py/objexcept.c index de9bf1694..f96523ed2 100644 --- a/py/objexcept.c +++ b/py/objexcept.c @@ -8,6 +8,8 @@ #include "qstr.h" #include "obj.h" #include "objtuple.h" +#include "runtime.h" +#include "runtime0.h" // This is unified class for C-level and Python-level exceptions // Python-level exceptions have empty ->msg and all arguments are in @@ -156,6 +158,11 @@ mp_obj_t mp_obj_new_exception(const mp_obj_type_t *exc_type) { return mp_obj_new_exception_msg_varg(exc_type, NULL); } +mp_obj_t mp_obj_new_exception_args(const mp_obj_type_t *exc_type, uint n_args, const mp_obj_t *args) { + assert(exc_type->make_new == mp_obj_exception_make_new); + return exc_type->make_new((mp_obj_t)exc_type, n_args, 0, args); +} + mp_obj_t mp_obj_new_exception_msg(const mp_obj_type_t *exc_type, const char *msg) { return mp_obj_new_exception_msg_varg(exc_type, msg); } @@ -202,6 +209,13 @@ bool mp_obj_is_exception_instance(mp_obj_t self_in) { return mp_obj_get_type(self_in)->make_new == mp_obj_exception_make_new; } +// return true if exception (type or instance) is a subclass of given +// exception type. +bool mp_obj_exception_match(mp_obj_t exc, const mp_obj_type_t *exc_type) { + // TODO: move implementation from RT_BINARY_OP_EXCEPTION_MATCH here. + return rt_binary_op(RT_BINARY_OP_EXCEPTION_MATCH, exc, (mp_obj_t)exc_type) == mp_const_true; +} + void mp_obj_exception_clear_traceback(mp_obj_t self_in) { // make sure self_in is an exception instance assert(mp_obj_get_type(self_in)->make_new == mp_obj_exception_make_new); diff --git a/py/objgenerator.c b/py/objgenerator.c index 39e362e04..3cfa02a0a 100644 --- a/py/objgenerator.c +++ b/py/objgenerator.c @@ -8,6 +8,7 @@ #include "obj.h" #include "runtime.h" #include "bc.h" +#include "objgenerator.h" /******************************************************************************/ /* generator wrapper */ @@ -73,9 +74,10 @@ mp_obj_t gen_instance_getiter(mp_obj_t self_in) { return self_in; } -STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value) { +mp_obj_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value, mp_vm_return_kind_t *ret_kind) { mp_obj_gen_instance_t *self = self_in; if (self->ip == 0) { + *ret_kind = MP_VM_RETURN_NORMAL; return mp_const_stop_iteration; } if (self->sp == self->state - 1) { @@ -85,10 +87,11 @@ STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw } else { *self->sp = send_value; } - mp_vm_return_kind_t vm_return_kind = mp_execute_byte_code_2(self->code_info, &self->ip, + *ret_kind = mp_execute_byte_code_2(self->code_info, &self->ip, &self->state[self->n_state - 1], &self->sp, (mp_exc_stack*)(self->state + self->n_state), &self->exc_sp, throw_value); - switch (vm_return_kind) { + + switch (*ret_kind) { case MP_VM_RETURN_NORMAL: // Explicitly mark generator as completed. If we don't do this, // subsequent next() may re-execute statements after last yield @@ -96,19 +99,39 @@ STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw // TODO: check how return with value behaves under such conditions // in CPython. self->ip = 0; - if (*self->sp == mp_const_none) { - return mp_const_stop_iteration; - } else { - // TODO return StopIteration with value *self->sp - return mp_const_stop_iteration; - } + return *self->sp; case MP_VM_RETURN_YIELD: return *self->sp; case MP_VM_RETURN_EXCEPTION: self->ip = 0; - nlr_jump(self->state[self->n_state - 1]); + return self->state[self->n_state - 1]; + + default: + assert(0); + return mp_const_none; + } +} + +STATIC mp_obj_t gen_resume_and_raise(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value) { + mp_vm_return_kind_t ret_kind; + mp_obj_t ret = mp_obj_gen_resume(self_in, send_value, throw_value, &ret_kind); + + switch (ret_kind) { + case MP_VM_RETURN_NORMAL: + // Optimize return w/o value in case generator is used in for loop + if (ret == mp_const_none) { + return mp_const_stop_iteration; + } else { + nlr_jump(mp_obj_new_exception_args(&mp_type_StopIteration, 1, &ret)); + } + + case MP_VM_RETURN_YIELD: + return ret; + + case MP_VM_RETURN_EXCEPTION: + nlr_jump(ret); default: assert(0); @@ -117,11 +140,11 @@ STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw } mp_obj_t gen_instance_iternext(mp_obj_t self_in) { - return gen_resume(self_in, mp_const_none, MP_OBJ_NULL); + return gen_resume_and_raise(self_in, mp_const_none, MP_OBJ_NULL); } STATIC mp_obj_t gen_instance_send(mp_obj_t self_in, mp_obj_t send_value) { - mp_obj_t ret = gen_resume(self_in, send_value, MP_OBJ_NULL); + mp_obj_t ret = gen_resume_and_raise(self_in, send_value, MP_OBJ_NULL); if (ret == mp_const_stop_iteration) { nlr_jump(mp_obj_new_exception(&mp_type_StopIteration)); } else { @@ -132,7 +155,7 @@ STATIC mp_obj_t gen_instance_send(mp_obj_t self_in, mp_obj_t send_value) { STATIC MP_DEFINE_CONST_FUN_OBJ_2(gen_instance_send_obj, gen_instance_send); STATIC mp_obj_t gen_instance_throw(uint n_args, const mp_obj_t *args) { - mp_obj_t ret = gen_resume(args[0], mp_const_none, n_args == 2 ? args[1] : args[2]); + mp_obj_t ret = gen_resume_and_raise(args[0], mp_const_none, n_args == 2 ? args[1] : args[2]); if (ret == mp_const_stop_iteration) { nlr_jump(mp_obj_new_exception(&mp_type_StopIteration)); } else { @@ -142,8 +165,30 @@ STATIC mp_obj_t gen_instance_throw(uint n_args, const mp_obj_t *args) { STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(gen_instance_throw_obj, 2, 4, gen_instance_throw); +STATIC mp_obj_t gen_instance_close(mp_obj_t self_in) { + mp_vm_return_kind_t ret_kind; + mp_obj_t ret = mp_obj_gen_resume(self_in, mp_const_none, (mp_obj_t)&mp_type_GeneratorExit, &ret_kind); + + if (ret_kind == MP_VM_RETURN_YIELD) { + nlr_jump(mp_obj_new_exception_msg(&mp_type_RuntimeError, "generator ignored GeneratorExit")); + } + // Swallow StopIteration & GeneratorExit (== successful close), and re-raise any other + if (ret_kind == MP_VM_RETURN_EXCEPTION) { + if (mp_obj_exception_match(ret, &mp_type_GeneratorExit) || + mp_obj_exception_match(ret, &mp_type_StopIteration)) { + return mp_const_none; + } + nlr_jump(ret); + } + + // The only choice left is MP_VM_RETURN_NORMAL which is successful close + return mp_const_none; +} + +STATIC MP_DEFINE_CONST_FUN_OBJ_1(gen_instance_close_obj, gen_instance_close); STATIC const mp_method_t gen_type_methods[] = { + { "close", &gen_instance_close_obj }, { "send", &gen_instance_send_obj }, { "throw", &gen_instance_throw_obj }, { NULL, NULL }, // end-of-list sentinel diff --git a/py/objgenerator.h b/py/objgenerator.h new file mode 100644 index 000000000..3dc69aa0e --- /dev/null +++ b/py/objgenerator.h @@ -0,0 +1 @@ +mp_obj_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_val, mp_obj_t throw_val, mp_vm_return_kind_t *ret_kind); diff --git a/tests/basics/generator-return.py b/tests/basics/generator-return.py new file mode 100644 index 000000000..a3ac88575 --- /dev/null +++ b/tests/basics/generator-return.py @@ -0,0 +1,10 @@ +def gen(): + yield 1 + return 42 + +g = gen() +print(next(g)) +try: + print(next(g)) +except StopIteration as e: + print(repr(e)) diff --git a/tests/basics/generator_close.py b/tests/basics/generator_close.py new file mode 100644 index 000000000..aa563f2a8 --- /dev/null +++ b/tests/basics/generator_close.py @@ -0,0 +1,59 @@ +def gen1(): + yield 1 + yield 2 + +# Test that it's possible to close just created gen +g = gen1() +print(g.close()) +try: + next(g) +except StopIteration: + print("StopIteration") + +# Test that it's possible to close gen in progress +g = gen1() +print(next(g)) +print(g.close()) +try: + next(g) + print("No StopIteration") +except StopIteration: + print("StopIteration") + +# Test that it's possible to close finished gen +g = gen1() +print(list(g)) +print(g.close()) +try: + next(g) + print("No StopIteration") +except StopIteration: + print("StopIteration") + + +# Throwing StopIteration in response to close() is ok +def gen2(): + try: + yield 1 + yield 2 + except: + raise StopIteration + +g = gen2() +next(g) +print(g.close()) + +# Any other exception is propagated to the .close() caller +def gen3(): + try: + yield 1 + yield 2 + except: + raise ValueError + +g = gen3() +next(g) +try: + print(g.close()) +except ValueError: + print("ValueError")