gh-116767: fix crash on 'async with' with many context managers (GH-118348)

Account for `add_stopiteration_handler` pushing a block for `async with`.
To allow generator functions that previously almost hit the `CO_MAXBLOCKS`
limit by nesting non-async blocks, the limit is increased by 1.
This increase allows one more block in non-generator functions.
This commit is contained in:
Irit Katriel 2024-05-01 12:01:16 +01:00 committed by GitHub
parent f6fab21721
commit c1bf4874c1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 50 additions and 7 deletions

View file

@ -113,7 +113,8 @@ compiler IR.
enum fblocktype { WHILE_LOOP, FOR_LOOP, TRY_EXCEPT, FINALLY_TRY, FINALLY_END,
WITH, ASYNC_WITH, HANDLER_CLEANUP, POP_VALUE, EXCEPTION_HANDLER,
EXCEPTION_GROUP_HANDLER, ASYNC_COMPREHENSION_GENERATOR };
EXCEPTION_GROUP_HANDLER, ASYNC_COMPREHENSION_GENERATOR,
STOP_ITERATION };
struct fblockinfo {
enum fblocktype fb_type;
@ -1503,6 +1504,7 @@ compiler_unwind_fblock(struct compiler *c, location *ploc,
case EXCEPTION_HANDLER:
case EXCEPTION_GROUP_HANDLER:
case ASYNC_COMPREHENSION_GENERATOR:
case STOP_ITERATION:
return SUCCESS;
case FOR_LOOP:
@ -2232,14 +2234,26 @@ compiler_function_body(struct compiler *c, stmt_ty s, int is_async, Py_ssize_t f
c->u->u_metadata.u_argcount = asdl_seq_LEN(args->args);
c->u->u_metadata.u_posonlyargcount = asdl_seq_LEN(args->posonlyargs);
c->u->u_metadata.u_kwonlyargcount = asdl_seq_LEN(args->kwonlyargs);
NEW_JUMP_TARGET_LABEL(c, start);
USE_LABEL(c, start);
bool add_stopiteration_handler = c->u->u_ste->ste_coroutine || c->u->u_ste->ste_generator;
if (add_stopiteration_handler) {
/* wrap_in_stopiteration_handler will push a block, so we need to account for that */
RETURN_IF_ERROR(
compiler_push_fblock(c, NO_LOCATION, STOP_ITERATION,
start, NO_LABEL, NULL));
}
for (Py_ssize_t i = first_instr; i < asdl_seq_LEN(body); i++) {
VISIT_IN_SCOPE(c, stmt, (stmt_ty)asdl_seq_GET(body, i));
}
if (c->u->u_ste->ste_coroutine || c->u->u_ste->ste_generator) {
if (add_stopiteration_handler) {
if (wrap_in_stopiteration_handler(c) < 0) {
compiler_exit_scope(c);
return ERROR;
}
compiler_pop_fblock(c, STOP_ITERATION, start);
}
PyCodeObject *co = optimize_and_assemble(c, 1);
compiler_exit_scope(c);