Merge pull request #228 from pfalcon/gen-send

Implement send() method for generators.

3 files changed, 35 insertions, 1 deletions

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))