Skip to content

Commit 962b1cd

Browse files
pfalconpfalcon
committed
objgenerator: Implement return with value and .close() method.
Return with value gets converted to StopIteration(value). Implementation keeps optimizing against creating of possibly unneeded exception objects, so there're considerable refactoring to implement these features.
1 parent 4b2b7ce commit 962b1cd

6 files changed

Lines changed: 143 additions & 12 deletions

File tree

‎py/obj.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,7 @@ mp_obj_t mp_obj_new_float(mp_float_t val);
273273
mp_obj_t mp_obj_new_complex(mp_float_t real, mp_float_t imag);
274274
#endif
275275
mp_obj_t mp_obj_new_exception(const mp_obj_type_t *exc_type);
276+
mp_obj_t mp_obj_new_exception_args(const mp_obj_type_t *exc_type, uint n_args, const mp_obj_t *args);
276277
mp_obj_t mp_obj_new_exception_msg(const mp_obj_type_t *exc_type, const char *msg);
277278
mp_obj_t mp_obj_new_exception_msg_varg(const mp_obj_type_t *exc_type, const char *fmt, ...); // counts args by number of % symbols in fmt, excluding %%; can only handle void* sizes (ie no float/double!)
278279
mp_obj_t mp_obj_new_range(int start, int stop, int step);
@@ -342,6 +343,7 @@ machine_int_t mp_obj_int_get_checked(mp_obj_t self_in);
342343
// exception
343344
bool mp_obj_is_exception_type(mp_obj_t self_in);
344345
bool mp_obj_is_exception_instance(mp_obj_t self_in);
346+
bool mp_obj_exception_match(mp_obj_t exc, const mp_obj_type_t *exc_type);
345347
void mp_obj_exception_clear_traceback(mp_obj_t self_in);
346348
void mp_obj_exception_add_traceback(mp_obj_t self_in, qstr file, machine_uint_t line, qstr block);
347349
void mp_obj_exception_get_traceback(mp_obj_t self_in, machine_uint_t *n, machine_uint_t **values);

‎py/objexcept.c‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
#include "qstr.h"
99
#include "obj.h"
1010
#include "objtuple.h"
11+
#include "runtime.h"
12+
#include "runtime0.h"
1113

1214
// This is unified class for C-level and Python-level exceptions
1315
// Python-level exceptions have empty ->msg and all arguments are in
@@ -156,6 +158,11 @@ mp_obj_t mp_obj_new_exception(const mp_obj_type_t *exc_type) {
156158
return mp_obj_new_exception_msg_varg(exc_type, NULL);
157159
}
158160

161+
mp_obj_t mp_obj_new_exception_args(const mp_obj_type_t *exc_type, uint n_args, const mp_obj_t *args) {
162+
assert(exc_type->make_new == mp_obj_exception_make_new);
163+
return exc_type->make_new((mp_obj_t)exc_type, n_args, 0, args);
164+
}
165+
159166
mp_obj_t mp_obj_new_exception_msg(const mp_obj_type_t *exc_type, const char *msg) {
160167
return mp_obj_new_exception_msg_varg(exc_type, msg);
161168
}
@@ -202,6 +209,13 @@ bool mp_obj_is_exception_instance(mp_obj_t self_in) {
202209
return mp_obj_get_type(self_in)->make_new == mp_obj_exception_make_new;
203210
}
204211

212+
// return true if exception (type or instance) is a subclass of given
213+
// exception type.
214+
bool mp_obj_exception_match(mp_obj_t exc, const mp_obj_type_t *exc_type) {
215+
// TODO: move implementation from RT_BINARY_OP_EXCEPTION_MATCH here.
216+
return rt_binary_op(RT_BINARY_OP_EXCEPTION_MATCH, exc, (mp_obj_t)exc_type) == mp_const_true;
217+
}
218+
205219
void mp_obj_exception_clear_traceback(mp_obj_t self_in) {
206220
// make sure self_in is an exception instance
207221
assert(mp_obj_get_type(self_in)->make_new == mp_obj_exception_make_new);

‎py/objgenerator.c‎

Lines changed: 57 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#include "obj.h"
99
#include "runtime.h"
1010
#include "bc.h"
11+
#include "objgenerator.h"
1112

1213
/******************************************************************************/
1314
/* generator wrapper */
@@ -73,9 +74,10 @@ mp_obj_t gen_instance_getiter(mp_obj_t self_in) {
7374
return self_in;
7475
}
7576

76-
STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value) {
77+
mp_obj_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value, mp_vm_return_kind_t *ret_kind) {
7778
mp_obj_gen_instance_t *self = self_in;
7879
if (self->ip == 0) {
80+
*ret_kind = MP_VM_RETURN_NORMAL;
7981
return mp_const_stop_iteration;
8082
}
8183
if (self->sp == self->state - 1) {
@@ -85,30 +87,51 @@ STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw
8587
} else {
8688
*self->sp = send_value;
8789
}
88-
mp_vm_return_kind_t vm_return_kind = mp_execute_byte_code_2(self->code_info, &self->ip,
90+
*ret_kind = mp_execute_byte_code_2(self->code_info, &self->ip,
8991
&self->state[self->n_state - 1], &self->sp, (mp_exc_stack*)(self->state + self->n_state),
9092
&self->exc_sp, throw_value);
91-
switch (vm_return_kind) {
93+
94+
switch (*ret_kind) {
9295
case MP_VM_RETURN_NORMAL:
9396
// Explicitly mark generator as completed. If we don't do this,
9497
// subsequent next() may re-execute statements after last yield
9598
// again and again, leading to side effects.
9699
// TODO: check how return with value behaves under such conditions
97100
// in CPython.
98101
self->ip = 0;
99-
if (*self->sp == mp_const_none) {
102+
return *self->sp;
103+
104+
case MP_VM_RETURN_YIELD:
105+
return *self->sp;
106+
107+
case MP_VM_RETURN_EXCEPTION:
108+
self->ip = 0;
109+
return self->state[self->n_state - 1];
110+
111+
default:
112+
assert(0);
113+
return mp_const_none;
114+
}
115+
}
116+
117+
STATIC mp_obj_t gen_resume_and_raise(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw_value) {
118+
mp_vm_return_kind_t ret_kind;
119+
mp_obj_t ret = mp_obj_gen_resume(self_in, send_value, throw_value, &ret_kind);
120+
121+
switch (ret_kind) {
122+
case MP_VM_RETURN_NORMAL:
123+
// Optimize return w/o value in case generator is used in for loop
124+
if (ret == mp_const_none) {
100125
return mp_const_stop_iteration;
101126
} else {
102-
// TODO return StopIteration with value *self->sp
103-
return mp_const_stop_iteration;
127+
nlr_jump(mp_obj_new_exception_args(&mp_type_StopIteration, 1, &ret));
104128
}
105129

106130
case MP_VM_RETURN_YIELD:
107-
return *self->sp;
131+
return ret;
108132

109133
case MP_VM_RETURN_EXCEPTION:
110-
self->ip = 0;
111-
nlr_jump(self->state[self->n_state - 1]);
134+
nlr_jump(ret);
112135

113136
default:
114137
assert(0);
@@ -117,11 +140,11 @@ STATIC mp_obj_t gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t throw
117140
}
118141

119142
mp_obj_t gen_instance_iternext(mp_obj_t self_in) {
120-
return gen_resume(self_in, mp_const_none, MP_OBJ_NULL);
143+
return gen_resume_and_raise(self_in, mp_const_none, MP_OBJ_NULL);
121144
}
122145

123146
STATIC mp_obj_t gen_instance_send(mp_obj_t self_in, mp_obj_t send_value) {
124-
mp_obj_t ret = gen_resume(self_in, send_value, MP_OBJ_NULL);
147+
mp_obj_t ret = gen_resume_and_raise(self_in, send_value, MP_OBJ_NULL);
125148
if (ret == mp_const_stop_iteration) {
126149
nlr_jump(mp_obj_new_exception(&mp_type_StopIteration));
127150
} else {
@@ -132,7 +155,7 @@ STATIC mp_obj_t gen_instance_send(mp_obj_t self_in, mp_obj_t send_value) {
132155
STATIC MP_DEFINE_CONST_FUN_OBJ_2(gen_instance_send_obj, gen_instance_send);
133156

134157
STATIC mp_obj_t gen_instance_throw(uint n_args, const mp_obj_t *args) {
135-
mp_obj_t ret = gen_resume(args[0], mp_const_none, n_args == 2 ? args[1] : args[2]);
158+
mp_obj_t ret = gen_resume_and_raise(args[0], mp_const_none, n_args == 2 ? args[1] : args[2]);
136159
if (ret == mp_const_stop_iteration) {
137160
nlr_jump(mp_obj_new_exception(&mp_type_StopIteration));
138161
} else {
@@ -142,8 +165,30 @@ STATIC mp_obj_t gen_instance_throw(uint n_args, const mp_obj_t *args) {
142165

143166
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(gen_instance_throw_obj, 2, 4, gen_instance_throw);
144167

168+
STATIC mp_obj_t gen_instance_close(mp_obj_t self_in) {
169+
mp_vm_return_kind_t ret_kind;
170+
mp_obj_t ret = mp_obj_gen_resume(self_in, mp_const_none, (mp_obj_t)&mp_type_GeneratorExit, &ret_kind);
171+
172+
if (ret_kind == MP_VM_RETURN_YIELD) {
173+
nlr_jump(mp_obj_new_exception_msg(&mp_type_RuntimeError, "generator ignored GeneratorExit"));
174+
}
175+
// Swallow StopIteration & GeneratorExit (== successful close), and re-raise any other
176+
if (ret_kind == MP_VM_RETURN_EXCEPTION) {
177+
if (mp_obj_exception_match(ret, &mp_type_GeneratorExit) ||
178+
mp_obj_exception_match(ret, &mp_type_StopIteration)) {
179+
return mp_const_none;
180+
}
181+
nlr_jump(ret);
182+
}
183+
184+
// The only choice left is MP_VM_RETURN_NORMAL which is successful close
185+
return mp_const_none;
186+
}
187+
188+
STATIC MP_DEFINE_CONST_FUN_OBJ_1(gen_instance_close_obj, gen_instance_close);
145189

146190
STATIC const mp_method_t gen_type_methods[] = {
191+
{ "close", &gen_instance_close_obj },
147192
{ "send", &gen_instance_send_obj },
148193
{ "throw", &gen_instance_throw_obj },
149194
{ NULL, NULL }, // end-of-list sentinel

‎py/objgenerator.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
mp_obj_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_val, mp_obj_t throw_val, mp_vm_return_kind_t *ret_kind);

‎tests/basics/generator-return.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
def gen():
2+
yield 1
3+
return 42
4+
5+
g = gen()
6+
print(next(g))
7+
try:
8+
print(next(g))
9+
except StopIteration as e:
10+
print(repr(e))

‎tests/basics/generator_close.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
def gen1():
2+
yield 1
3+
yield 2
4+
5+
# Test that it's possible to close just created gen
6+
g = gen1()
7+
print(g.close())
8+
try:
9+
next(g)
10+
except StopIteration:
11+
print("StopIteration")
12+
13+
# Test that it's possible to close gen in progress
14+
g = gen1()
15+
print(next(g))
16+
print(g.close())
17+
try:
18+
next(g)
19+
print("No StopIteration")
20+
except StopIteration:
21+
print("StopIteration")
22+
23+
# Test that it's possible to close finished gen
24+
g = gen1()
25+
print(list(g))
26+
print(g.close())
27+
try:
28+
next(g)
29+
print("No StopIteration")
30+
except StopIteration:
31+
print("StopIteration")
32+
33+
34+
# Throwing StopIteration in response to close() is ok
35+
def gen2():
36+
try:
37+
yield 1
38+
yield 2
39+
except:
40+
raise StopIteration
41+
42+
g = gen2()
43+
next(g)
44+
print(g.close())
45+
46+
# Any other exception is propagated to the .close() caller
47+
def gen3():
48+
try:
49+
yield 1
50+
yield 2
51+
except:
52+
raise ValueError
53+
54+
g = gen3()
55+
next(g)
56+
try:
57+
print(g.close())
58+
except ValueError:
59+
print("ValueError")

0 commit comments

Comments
 (0)