summaryrefslogtreecommitdiffstatshomepage
path: root/py
diff options
context:
space:
mode:
Diffstat (limited to 'py')
-rw-r--r--py/objgenerator.c20
-rw-r--r--py/vm.c1
2 files changed, 20 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)