diff --git a/py/objgenerator.c b/py/objgenerator.c index 192e5c6328..0cac34b09e 100644 --- a/py/objgenerator.c +++ b/py/objgenerator.c @@ -73,8 +73,15 @@ mp_obj_t gen_instance_getiter(mp_obj_t self_in) { return self_in; } -mp_obj_t gen_instance_iternext(mp_obj_t self_in) { +static mp_obj_t gen_send(mp_obj_t self_in, mp_obj_t send_value) { mp_obj_gen_instance_t *self = self_in; + if (self->sp == self->state - 1) { + if (send_value != mp_const_none) { + nlr_jump(mp_obj_new_exception_msg(MP_QSTR_TypeError, "can't send non-None value to a just-started generator")); + } + } else { + *self->sp = send_value; + } bool yield = mp_execute_byte_code_2(self->code_info, &self->ip, &self->state[self->n_state - 1], &self->sp); if (yield) { return *self->sp; @@ -87,6 +94,16 @@ mp_obj_t gen_instance_iternext(mp_obj_t self_in) { } } } +static MP_DEFINE_CONST_FUN_OBJ_2(gen_send_obj, gen_send); + +mp_obj_t gen_instance_iternext(mp_obj_t self_in) { + return gen_send(self_in, mp_const_none); +} + +static const mp_method_t gen_type_methods[] = { + { "send", &gen_send_obj }, + { NULL, NULL }, // end-of-list sentinel +}; const mp_obj_type_t gen_instance_type = { { &mp_const_type }, @@ -94,6 +111,7 @@ const mp_obj_type_t gen_instance_type = { .print = gen_instance_print, .getiter = gen_instance_getiter, .iternext = gen_instance_iternext, + .methods = gen_type_methods, }; mp_obj_t mp_obj_new_gen_instance(const byte *bytecode, uint n_state, int n_args, const mp_obj_t *args) { diff --git a/py/vm.c b/py/vm.c index 82a9f893f3..0cc26021e8 100644 --- a/py/vm.c +++ b/py/vm.c @@ -73,6 +73,7 @@ mp_obj_t mp_execute_byte_code(const byte *code, const mp_obj_t *args, uint n_arg // fastn has items in reverse order (fastn[0] is local[0], fastn[-1] is local[1], etc) // sp points to bottom of stack which grows up +// returns true if bytecode yielded bool mp_execute_byte_code_2(const byte *code_info, const byte **ip_in_out, mp_obj_t *fastn, mp_obj_t **sp_in_out) { // careful: be sure to declare volatile any variables read in the exception handler (written is ok, I think) diff --git a/tests/basics/generator_send.py b/tests/basics/generator_send.py new file mode 100644 index 0000000000..4158478cac --- /dev/null +++ b/tests/basics/generator_send.py @@ -0,0 +1,15 @@ +def f(): + n = 0 + while True: + n = yield n + 1 + print(n) + +g = f() +try: + g.send(1) +except TypeError: + print("caught") + +print(g.send(None)) +print(g.send(100)) +print(g.send(200))