Skip to content

Commit 9cda4fc

Browse files
markshannonmarkshannon
committed
Wrap async generators and generator expressions in handler for Stop[Async]Iteration.
1 parent 65d303f commit 9cda4fc

File tree

5 files changed

+80
-57
lines changed

5 files changed

+80
-57
lines changed

‎Lib/test/test_dis.py‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -545,10 +545,12 @@ async def _asyncwith(c):
545545
RERAISE 1
546546
>> LOAD_ERROR 1
547547
CHECK_EXC_MATCH
548-
POP_JUMP_IF_FALSE 8 (to 140)
548+
POP_JUMP_IF_FALSE 10 (to 144)
549+
PUSH_EXC_INFO
549550
LOAD_ERROR 2
550551
LOAD_CONST 3 ('coroutine raised StopIteration')
551552
CALL 0
553+
SWAP 2
552554
RAISE_VARARGS 2
553555
>> RERAISE 1
554556
ExceptionTable:
@@ -710,6 +712,18 @@ def foo(x):
710712
JUMP_BACKWARD 9 (to 8)
711713
>> END_FOR
712714
RETURN_VALUE
715+
>> LOAD_ERROR 1
716+
CHECK_EXC_MATCH
717+
POP_JUMP_IF_FALSE 10 (to 56)
718+
PUSH_EXC_INFO
719+
LOAD_ERROR 2
720+
LOAD_CONST 1 ('generator raised StopIteration')
721+
CALL 0
722+
SWAP 2
723+
RAISE_VARARGS 2
724+
>> RERAISE 1
725+
ExceptionTable:
726+
1 row
713727
""" % (dis_nested_1,
714728
__file__,
715729
_h.__code__.co_firstlineno + 3,

‎Lib/test/test_sys.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1439,7 +1439,7 @@ def bar(cls):
14391439
check(bar, size('PP'))
14401440
# generator
14411441
def get_gen(): yield 1
1442-
check(get_gen(), size('P2P4P4c7P2ic??4P'))
1442+
check(get_gen(), size('P2P4P4c7P2ic??5P'))
14431443
# iterator
14441444
check(iter('abc'), size('lP'))
14451445
# callable-iterator

‎Objects/genobject.c‎

Lines changed: 3 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -247,25 +247,9 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
247247
}
248248
}
249249
else {
250-
if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
251-
const char *msg = "generator raised StopIteration";
252-
if (PyCoro_CheckExact(gen)) {
253-
msg = "coroutine raised StopIteration";
254-
}
255-
else if (PyAsyncGen_CheckExact(gen)) {
256-
msg = "async generator raised StopIteration";
257-
}
258-
_PyErr_FormatFromCause(PyExc_RuntimeError, "%s", msg);
259-
}
260-
else if (PyAsyncGen_CheckExact(gen) &&
261-
PyErr_ExceptionMatches(PyExc_StopAsyncIteration))
262-
{
263-
/* code in `gen` raised a StopAsyncIteration error:
264-
raise a RuntimeError.
265-
*/
266-
const char *msg = "async generator raised StopAsyncIteration";
267-
_PyErr_FormatFromCause(PyExc_RuntimeError, "%s", msg);
268-
}
250+
assert(!PyErr_ExceptionMatches(PyExc_StopIteration));
251+
assert(!PyAsyncGen_CheckExact(gen) ||
252+
!PyErr_ExceptionMatches(PyExc_StopAsyncIteration));
269253
}
270254

271255
/* generator can't be rerun, so release the frame */

‎Python/ceval.c‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2187,6 +2187,9 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
21872187
case 2:
21882188
value = PyExc_RuntimeError;
21892189
break;
2190+
case 3:
2191+
value = PyExc_StopAsyncIteration;
2192+
break;
21902193
default:
21912194
Py_UNREACHABLE();
21922195
}

‎Python/compile.c‎

Lines changed: 58 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -2568,22 +2568,42 @@ compiler_check_debug_args(struct compiler *c, arguments_ty args)
25682568
return 1;
25692569
}
25702570

2571+
static inline int
2572+
insert_instruction(basicblock *block, int pos, struct instr *instr) {
2573+
if (basicblock_next_instr(block) < 0) {
2574+
return -1;
2575+
}
2576+
for (int i = block->b_iused - 1; i > pos; i--) {
2577+
block->b_instr[i] = block->b_instr[i-1];
2578+
}
2579+
block->b_instr[pos] = *instr;
2580+
return 0;
2581+
}
2582+
25712583
static int
2572-
coroutine_stopiteration_handler(struct compiler *c, jump_target_label *handler)
2584+
wrap_in_stopiteration_handler(struct compiler *c)
25732585
{
2574-
ADDOP_LOAD_CONST(c, NO_LOCATION, Py_None);
2575-
ADDOP(c, NO_LOCATION, RETURN_VALUE);
2576-
if (cfg_builder_use_label(CFG_BUILDER(c), *handler) < 0) {
2577-
return 0;
2578-
}
2579-
jump_target_label other = cfg_new_label(CFG_BUILDER(c));
2580-
if (!IS_LABEL(other)) {
2586+
NEW_JUMP_TARGET_LABEL(c, handler);
2587+
NEW_JUMP_TARGET_LABEL(c, next);
2588+
2589+
/* Insert SETUP_CLEANUP at start */
2590+
struct instr setup = {
2591+
.i_opcode = SETUP_CLEANUP,
2592+
.i_oparg = handler.id,
2593+
.i_loc = NO_LOCATION,
2594+
.i_target = NULL,
2595+
};
2596+
if (insert_instruction(c->u->u_cfg_builder.g_entryblock, 0, &setup)) {
25812597
return 0;
25822598
}
2583-
USE_LABEL(c, *handler);
2599+
2600+
ADDOP_LOAD_CONST(c, NO_LOCATION, Py_None);
2601+
ADDOP(c, NO_LOCATION, RETURN_VALUE);
2602+
USE_LABEL(c, handler);
25842603
ADDOP_I(c, NO_LOCATION, LOAD_ERROR, 1); // StopIteration
25852604
ADDOP(c, NO_LOCATION, CHECK_EXC_MATCH);
2586-
ADDOP_JUMP(c, NO_LOCATION, POP_JUMP_IF_FALSE, other);
2605+
ADDOP_JUMP(c, NO_LOCATION, POP_JUMP_IF_FALSE, next);
2606+
ADDOP(c, NO_LOCATION, PUSH_EXC_INFO);
25872607
ADDOP_I(c, NO_LOCATION, LOAD_ERROR, 2); // RuntimeError
25882608
const char *msg = c->u->u_ste->ste_coroutine ?
25892609
(c->u->u_ste->ste_generator ?
@@ -2597,9 +2617,31 @@ coroutine_stopiteration_handler(struct compiler *c, jump_target_label *handler)
25972617
}
25982618
ADDOP_LOAD_CONST_NEW(c, NO_LOCATION, message);
25992619
ADDOP_I(c, NO_LOCATION, CALL, 0);
2620+
ADDOP_I(c, NO_LOCATION, SWAP, 2);
26002621
ADDOP_I(c, NO_LOCATION, RAISE_VARARGS, 2);
26012622

2602-
USE_LABEL(c, other);
2623+
USE_LABEL(c, next);
2624+
if (c->u->u_ste->ste_coroutine &&
2625+
c->u->u_ste->ste_generator)
2626+
{
2627+
NEW_JUMP_TARGET_LABEL(c, next);
2628+
ADDOP_I(c, NO_LOCATION, LOAD_ERROR, 3); // StopAsyncIteration
2629+
ADDOP(c, NO_LOCATION, CHECK_EXC_MATCH);
2630+
ADDOP_JUMP(c, NO_LOCATION, POP_JUMP_IF_FALSE, next);
2631+
ADDOP(c, NO_LOCATION, PUSH_EXC_INFO);
2632+
ADDOP_I(c, NO_LOCATION, LOAD_ERROR, 2); // RuntimeError
2633+
const char *msg = "async generator raised StopAsyncIteration";
2634+
PyObject *message = _PyUnicode_FromASCII(msg, strlen(msg));
2635+
if (message == NULL) {
2636+
return 0;
2637+
}
2638+
ADDOP_LOAD_CONST_NEW(c, NO_LOCATION, message);
2639+
ADDOP_I(c, NO_LOCATION, CALL, 0);
2640+
ADDOP_I(c, NO_LOCATION, SWAP, 2);
2641+
ADDOP_I(c, NO_LOCATION, RAISE_VARARGS, 2);
2642+
2643+
USE_LABEL(c, next);
2644+
}
26032645
ADDOP_I(c, NO_LOCATION, RERAISE, 1);
26042646
return 1;
26052647
}
@@ -2669,17 +2711,6 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async)
26692711
return 0;
26702712
}
26712713

2672-
bool is_coro = (c->u->u_ste->ste_coroutine || c->u->u_ste->ste_generator);
2673-
jump_target_label handler;
2674-
if (is_coro) {
2675-
handler = cfg_new_label(CFG_BUILDER(c));
2676-
if (!IS_LABEL(handler)) {
2677-
compiler_exit_scope(c);
2678-
return 0;
2679-
}
2680-
ADDOP_JUMP(c, NO_LOCATION, SETUP_CLEANUP, handler);
2681-
}
2682-
26832714
/* if not -OO mode, add docstring */
26842715
if (c->c_optimize < 2) {
26852716
docstring = _PyAST_GetDocString(body);
@@ -2695,8 +2726,8 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async)
26952726
for (i = docstring ? 1 : 0; i < asdl_seq_LEN(body); i++) {
26962727
VISIT_IN_SCOPE(c, stmt, (stmt_ty)asdl_seq_GET(body, i));
26972728
}
2698-
if (is_coro) {
2699-
if (!coroutine_stopiteration_handler(c, &handler)) {
2729+
if (c->u->u_ste->ste_coroutine || c->u->u_ste->ste_generator) {
2730+
if (!wrap_in_stopiteration_handler(c)) {
27002731
compiler_exit_scope(c);
27012732
return 0;
27022733
}
@@ -5542,6 +5573,9 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
55425573
if (type != COMP_GENEXP) {
55435574
ADDOP(c, LOC(e), RETURN_VALUE);
55445575
}
5576+
if (!wrap_in_stopiteration_handler(c)) {
5577+
goto error_in_scope;
5578+
}
55455579

55465580
co = assemble(c, 1);
55475581
qualname = c->u->u_qualname;
@@ -8606,18 +8640,6 @@ build_cellfixedoffsets(struct compiler *c)
86068640
return fixed;
86078641
}
86088642

8609-
static inline int
8610-
insert_instruction(basicblock *block, int pos, struct instr *instr) {
8611-
if (basicblock_next_instr(block) < 0) {
8612-
return -1;
8613-
}
8614-
for (int i = block->b_iused - 1; i > pos; i--) {
8615-
block->b_instr[i] = block->b_instr[i-1];
8616-
}
8617-
block->b_instr[pos] = *instr;
8618-
return 0;
8619-
}
8620-
86218643
static int
86228644
insert_prefix_instructions(struct compiler *c, basicblock *entryblock,
86238645
int *fixed, int nfreevars, int code_flags)

0 commit comments

Comments
 (0)