Skip to content

Commit da51a39

Browse files
dpgeorgedpgeorge
committed
Merge pull request adafruit#383 from pfalcon/yield-from
Implement "yield from"
2 parents 75f7158 + 3c2b2ac commit da51a39

6 files changed

Lines changed: 232 additions & 1 deletion

File tree

‎py/vm.c‎

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#include "runtime.h"
1010
#include "bc0.h"
1111
#include "bc.h"
12+
#include "objgenerator.h"
1213

1314
// Value stack grows up (this makes it incompatible with native C stack, but
1415
// makes sure that arguments to functions are in natural order arg1..argN
@@ -146,7 +147,9 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
146147
// If we have exception to inject, now that we finish setting up
147148
// execution context, raise it. This works as if RAISE_VARARGS
148149
// bytecode was executed.
149-
if (inject_exc != MP_OBJ_NULL) {
150+
// Injecting exc into yield from generator is a special case,
151+
// handled by MP_BC_YIELD_FROM itself
152+
if (inject_exc != MP_OBJ_NULL && *ip != MP_BC_YIELD_FROM) {
150153
mp_obj_t t = inject_exc;
151154
inject_exc = MP_OBJ_NULL;
152155
nlr_jump(rt_make_raise_obj(t));
@@ -634,12 +637,65 @@ mp_vm_return_kind_t mp_execute_byte_code_2(const byte *code_info, const byte **i
634637
nlr_jump(rt_make_raise_obj(obj1));
635638

636639
case MP_BC_YIELD_VALUE:
640+
yield:
637641
nlr_pop();
638642
*ip_in_out = ip;
639643
*sp_in_out = sp;
640644
*exc_sp_in_out = MP_TAGPTR_MAKE(exc_sp, currently_in_except_block);
641645
return MP_VM_RETURN_YIELD;
642646

647+
case MP_BC_YIELD_FROM:
648+
{
649+
//#define EXC_MATCH(exc, type) MP_OBJ_IS_TYPE(exc, type)
650+
#define EXC_MATCH(exc, type) mp_obj_exception_match(exc, type)
651+
#define GENERATOR_EXIT_IF_NEEDED(t) if (t != MP_OBJ_NULL && EXC_MATCH(t, &mp_type_GeneratorExit)) { nlr_jump(t); }
652+
mp_vm_return_kind_t ret_kind;
653+
obj1 = POP();
654+
mp_obj_t t_exc = MP_OBJ_NULL;
655+
if (inject_exc != MP_OBJ_NULL) {
656+
t_exc = inject_exc;
657+
inject_exc = MP_OBJ_NULL;
658+
ret_kind = mp_obj_gen_resume(TOP(), mp_const_none, t_exc, &obj2);
659+
} else {
660+
ret_kind = mp_obj_gen_resume(TOP(), obj1, MP_OBJ_NULL, &obj2);
661+
}
662+
663+
if (ret_kind == MP_VM_RETURN_YIELD) {
664+
ip--;
665+
PUSH(obj2);
666+
goto yield;
667+
}
668+
if (ret_kind == MP_VM_RETURN_NORMAL) {
669+
// Pop exhausted gen
670+
sp--;
671+
if (obj2 == MP_OBJ_NULL) {
672+
// Optimize StopIteration
673+
// TODO: get StopIteration's value
674+
PUSH(mp_const_none);
675+
} else {
676+
PUSH(obj2);
677+
}
678+
679+
// If we injected GeneratorExit downstream, then even
680+
// if it was swallowed, we re-raise GeneratorExit
681+
GENERATOR_EXIT_IF_NEEDED(t_exc);
682+
break;
683+
}
684+
if (ret_kind == MP_VM_RETURN_EXCEPTION) {
685+
// Pop exhausted gen
686+
sp--;
687+
if (EXC_MATCH(obj2, &mp_type_StopIteration)) {
688+
PUSH(mp_obj_exception_get_value(obj2));
689+
// If we injected GeneratorExit downstream, then even
690+
// if it was swallowed, we re-raise GeneratorExit
691+
GENERATOR_EXIT_IF_NEEDED(t_exc);
692+
break;
693+
} else {
694+
nlr_jump(obj2);
695+
}
696+
}
697+
}
698+
643699
case MP_BC_IMPORT_NAME:
644700
DECODE_QSTR;
645701
obj1 = POP();

‎tests/basics/gen-yield-from-close.py‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
def gen():
2+
yield 1
3+
yield 2
4+
yield 3
5+
yield 4
6+
7+
def gen2():
8+
yield -1
9+
print((yield from gen()))
10+
yield 10
11+
yield 11
12+
13+
g = gen2()
14+
print(next(g))
15+
print(next(g))
16+
g.close()
17+
try:
18+
print(next(g))
19+
except StopIteration:
20+
print("StopIteration")
21+
22+
23+
# Now variation of same test, but with leaf generator
24+
# swallowing GeneratorExit exception - its upstream gen
25+
# generator should still receive one.
26+
def gen3():
27+
yield 1
28+
try:
29+
yield 2
30+
except GeneratorExit:
31+
print("leaf caught GeneratorExit and swallowed it")
32+
return
33+
yield 3
34+
yield 4
35+
36+
def gen4():
37+
yield -1
38+
try:
39+
print((yield from gen3()))
40+
except GeneratorExit:
41+
print("delegating caught GeneratorExit")
42+
raise
43+
yield 10
44+
yield 11
45+
46+
g = gen4()
47+
print(next(g))
48+
print(next(g))
49+
print(next(g))
50+
g.close()
51+
try:
52+
print(next(g))
53+
except StopIteration:
54+
print("StopIteration")
55+
56+
57+
# Yet another variation - leaf generator gets GeneratorExit,
58+
# but raises StopIteration instead. This still should close chain properly.
59+
def gen5():
60+
yield 1
61+
try:
62+
yield 2
63+
except GeneratorExit:
64+
print("leaf caught GeneratorExit and raised StopIteration instead")
65+
raise StopIteration(123)
66+
yield 3
67+
yield 4
68+
69+
def gen6():
70+
yield -1
71+
try:
72+
print((yield from gen5()))
73+
except GeneratorExit:
74+
print("delegating caught GeneratorExit")
75+
raise
76+
yield 10
77+
yield 11
78+
79+
g = gen6()
80+
print(next(g))
81+
print(next(g))
82+
print(next(g))
83+
g.close()
84+
try:
85+
print(next(g))
86+
except StopIteration:
87+
print("StopIteration")

‎tests/basics/gen-yield-from-exc.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
def gen():
2+
yield 1
3+
yield 2
4+
raise ValueError
5+
6+
def gen2():
7+
try:
8+
print((yield from gen()))
9+
except ValueError:
10+
print("caught ValueError from downstream")
11+
12+
g = gen2()
13+
print(list(g))

‎tests/basics/gen-yield-from-send.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
def gen():
2+
print("sent:", (yield 1))
3+
yield 2
4+
5+
def gen2():
6+
print((yield from gen()))
7+
8+
g = gen2()
9+
next(g)
10+
print("yielded:", g.send("val"))
11+
try:
12+
next(g)
13+
except StopIteration:
14+
print("StopIteration")

‎tests/basics/gen-yield-from-throw.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
def gen():
2+
try:
3+
yield 1
4+
except ValueError:
5+
print("got ValueError from upstream!")
6+
yield "str1"
7+
raise TypeError
8+
9+
def gen2():
10+
print((yield from gen()))
11+
12+
g = gen2()
13+
print(next(g))
14+
print(g.throw(ValueError))
15+
try:
16+
print(next(g))
17+
except TypeError:
18+
print("got TypeError from downstream!")
19+

‎tests/basics/gen-yield-from.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# Case of terminating subgen using return with value
2+
def gen():
3+
yield 1
4+
yield 2
5+
return 3
6+
7+
def gen2():
8+
print("here1")
9+
print((yield from gen()))
10+
print("here2")
11+
12+
g = gen2()
13+
print(list(g))
14+
15+
16+
# Like above, but terminate subgen using StopIteration
17+
def gen3():
18+
yield 1
19+
yield 2
20+
raise StopIteration
21+
22+
def gen4():
23+
print("here1")
24+
print((yield from gen3()))
25+
print("here2")
26+
27+
g = gen4()
28+
print(list(g))
29+
30+
# Like above, but terminate subgen using StopIteration with value
31+
def gen5():
32+
yield 1
33+
yield 2
34+
raise StopIteration(123)
35+
36+
def gen6():
37+
print("here1")
38+
print((yield from gen5()))
39+
print("here2")
40+
41+
g = gen6()
42+
print(list(g))

0 commit comments

Comments
 (0)