py: Generators can have their locals closed over.

dpgeorge · dpgeorge · commit 13d6739cc7a5 · 2014-04-09T19:01:45.000+01:00
diff --git 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
@@ -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)
