gh-143055: Implementation of PEP 798 (#143056)

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
This commit is contained in:
adam j hartz 2026-01-30 23:37:52 -05:00 committed by GitHub
parent 96e4cd698a
commit ccbe41e27c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 4069 additions and 2908 deletions

View file

@ -4522,28 +4522,63 @@ codegen_sync_comprehension_generator(compiler *c, location loc,
/* comprehension specific code */
switch (type) {
case COMP_GENEXP:
VISIT(c, expr, elt);
ADDOP_YIELD(c, elt_loc);
ADDOP(c, elt_loc, POP_TOP);
if (elt->kind == Starred_kind) {
NEW_JUMP_TARGET_LABEL(c, unpack_start);
NEW_JUMP_TARGET_LABEL(c, unpack_end);
VISIT(c, expr, elt->v.Starred.value);
ADDOP(c, elt_loc, GET_ITER);
USE_LABEL(c, unpack_start);
ADDOP_JUMP(c, elt_loc, FOR_ITER, unpack_end);
ADDOP_YIELD(c, elt_loc);
ADDOP(c, elt_loc, POP_TOP);
ADDOP_JUMP(c, NO_LOCATION, JUMP, unpack_start);
USE_LABEL(c, unpack_end);
ADDOP(c, NO_LOCATION, END_FOR);
ADDOP(c, NO_LOCATION, POP_ITER);
}
else {
VISIT(c, expr, elt);
ADDOP_YIELD(c, elt_loc);
ADDOP(c, elt_loc, POP_TOP);
}
break;
case COMP_LISTCOMP:
VISIT(c, expr, elt);
ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
if (elt->kind == Starred_kind) {
VISIT(c, expr, elt->v.Starred.value);
ADDOP_I(c, elt_loc, LIST_EXTEND, depth + 1);
}
else {
VISIT(c, expr, elt);
ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
}
break;
case COMP_SETCOMP:
VISIT(c, expr, elt);
ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
if (elt->kind == Starred_kind) {
VISIT(c, expr, elt->v.Starred.value);
ADDOP_I(c, elt_loc, SET_UPDATE, depth + 1);
}
else {
VISIT(c, expr, elt);
ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
}
break;
case COMP_DICTCOMP:
/* With '{k: v}', k is evaluated before v, so we do
the same. */
VISIT(c, expr, elt);
VISIT(c, expr, val);
elt_loc = LOCATION(elt->lineno,
val->end_lineno,
elt->col_offset,
val->end_col_offset);
ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
if (val == NULL) {
/* unpacking (**) case */
VISIT(c, expr, elt);
ADDOP_I(c, elt_loc, DICT_UPDATE, depth+1);
}
else {
/* With '{k: v}', k is evaluated before v, so we do
the same. */
VISIT(c, expr, elt);
VISIT(c, expr, val);
elt_loc = LOCATION(elt->lineno,
val->end_lineno,
elt->col_offset,
val->end_col_offset);
ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
}
break;
default:
return ERROR;
@ -4629,28 +4664,63 @@ codegen_async_comprehension_generator(compiler *c, location loc,
/* comprehension specific code */
switch (type) {
case COMP_GENEXP:
VISIT(c, expr, elt);
ADDOP_YIELD(c, elt_loc);
ADDOP(c, elt_loc, POP_TOP);
if (elt->kind == Starred_kind) {
NEW_JUMP_TARGET_LABEL(c, unpack_start);
NEW_JUMP_TARGET_LABEL(c, unpack_end);
VISIT(c, expr, elt->v.Starred.value);
ADDOP(c, elt_loc, GET_ITER);
USE_LABEL(c, unpack_start);
ADDOP_JUMP(c, elt_loc, FOR_ITER, unpack_end);
ADDOP_YIELD(c, elt_loc);
ADDOP(c, elt_loc, POP_TOP);
ADDOP_JUMP(c, NO_LOCATION, JUMP, unpack_start);
USE_LABEL(c, unpack_end);
ADDOP(c, NO_LOCATION, END_FOR);
ADDOP(c, NO_LOCATION, POP_ITER);
}
else {
VISIT(c, expr, elt);
ADDOP_YIELD(c, elt_loc);
ADDOP(c, elt_loc, POP_TOP);
}
break;
case COMP_LISTCOMP:
VISIT(c, expr, elt);
ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
if (elt->kind == Starred_kind) {
VISIT(c, expr, elt->v.Starred.value);
ADDOP_I(c, elt_loc, LIST_EXTEND, depth + 1);
}
else {
VISIT(c, expr, elt);
ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
}
break;
case COMP_SETCOMP:
VISIT(c, expr, elt);
ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
if (elt->kind == Starred_kind) {
VISIT(c, expr, elt->v.Starred.value);
ADDOP_I(c, elt_loc, SET_UPDATE, depth + 1);
}
else {
VISIT(c, expr, elt);
ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
}
break;
case COMP_DICTCOMP:
/* With '{k: v}', k is evaluated before v, so we do
the same. */
VISIT(c, expr, elt);
VISIT(c, expr, val);
elt_loc = LOCATION(elt->lineno,
val->end_lineno,
elt->col_offset,
val->end_col_offset);
ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
if (val == NULL) {
/* unpacking (**) case */
VISIT(c, expr, elt);
ADDOP_I(c, elt_loc, DICT_UPDATE, depth+1);
}
else {
/* With '{k: v}', k is evaluated before v, so we do
the same. */
VISIT(c, expr, elt);
VISIT(c, expr, val);
elt_loc = LOCATION(elt->lineno,
val->end_lineno,
elt->col_offset,
val->end_col_offset);
ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
}
break;
default:
return ERROR;