diff options
| -rw-r--r-- | py/objgenerator.c | 22 | ||||
| -rw-r--r-- | tests/basics/generator-closure.py | 26 |
2 files changed, 40 insertions, 8 deletions
diff --git a/py/objgenerator.c b/py/objgenerator.c index 975681af88..3a7f9f9d40 100644 --- a/py/objgenerator.c +++ b/py/objgenerator.c @@ -236,17 +236,10 @@ mp_obj_t mp_obj_new_gen_instance(const byte *bytecode, uint n_args, const mp_obj machine_uint_t n_exc_stack = bytecode[2] | (bytecode[3] << 8); bytecode += 4; - // bytecode prelude: initialise closed over variables - // TODO - // for now we just make sure there are no cells variables - // need to work out how to implement closed over variables in generators - assert(bytecode[0] == 0); - bytecode += 1; - + // allocate the generator object, with room for local stack and exception stack mp_obj_gen_instance_t *o = m_new_obj_var(mp_obj_gen_instance_t, byte, n_state * sizeof(mp_obj_t) + n_exc_stack * sizeof(mp_exc_stack_t)); o->base.type = &mp_type_gen_instance; o->code_info = code_info; - o->ip = bytecode; o->sp = &o->state[0] - 1; // sp points to top of stack, which starts off 1 below the state o->exc_sp = (mp_exc_stack_t*)(o->state + n_state) - 1; o->n_state = n_state; @@ -259,5 +252,18 @@ mp_obj_t mp_obj_new_gen_instance(const byte *bytecode, uint n_args, const mp_obj o->state[n_state - 1 - n_args - i] = args2[i]; } + // bytecode prelude: initialise closed over variables + for (uint n_local = *bytecode++; n_local > 0; n_local--) { + uint local_num = *bytecode++; + if (local_num < n_args + n_args2) { + o->state[n_state - 1 - local_num] = mp_obj_new_cell(o->state[n_state - 1 - local_num]); + } else { + o->state[n_state - 1 - local_num] = mp_obj_new_cell(MP_OBJ_NULL); + } + } + + // set ip to start of actual byte code + o->ip = bytecode; + return o; } diff --git a/tests/basics/generator-closure.py b/tests/basics/generator-closure.py new file mode 100644 index 0000000000..d8a517edeb --- /dev/null +++ b/tests/basics/generator-closure.py @@ -0,0 +1,26 @@ +# a generator that closes over outer variables +def f(): + x = 1 # closed over by g + def g(): + yield x + yield x + 1 + return g() +for i in f(): + print(i) + +# a generator that has its variables closed over +def f(): + x = 1 # closed over by g + def g(): + return x + 1 + yield g() + x = 2 + yield g() +for i in f(): + print(i) + +# using comprehensions, the inner generator closes over y +generator_of_generators = (((x, y) for x in range(2)) for y in range(3)) +for i in generator_of_generators: + for j in i: + print(j) |