Commit bf38e2a0 authored by Paul Sokolovsky's avatar Paul Sokolovsky
Browse files

Implement send() method for generators.

parent 56bb6360
......@@ -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) {
......
......@@ -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)
......
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))
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment