Implement send() method for generators.

Author: pfalcon

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,13 +94,24 @@ 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 },
  "generator",
  .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) {

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)
 

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