Skip to content

Commit fef8ef4

Browse files
cdce8pcdce8p
committed
gh-151907: Reduce duplication in codegen comprehension functions
1 parent 1ec5607 commit fef8ef4

1 file changed

Lines changed: 88 additions & 136 deletions

File tree

‎Python/codegen.c‎

Lines changed: 88 additions & 136 deletions
Original file line numberDiff line numberDiff line change
@@ -4586,6 +4586,78 @@ codegen_unpack_starred(compiler *c, location loc, expr_ty value, bool yield)
45864586
return SUCCESS;
45874587
}
45884588

4589+
static int
4590+
codegen_comprehension_generator_helper(compiler *c, location elt_loc, int depth,
4591+
expr_ty elt, expr_ty val, int type,
4592+
bool avoid_creation)
4593+
{
4594+
switch (type) {
4595+
case COMP_GENEXP:
4596+
assert(!avoid_creation);
4597+
if (elt->kind == Starred_kind) {
4598+
RETURN_IF_ERROR(codegen_unpack_starred(c, elt_loc, elt->v.Starred.value, /*yield=*/true));
4599+
}
4600+
else {
4601+
VISIT(c, expr, elt);
4602+
ADDOP_YIELD(c, elt_loc);
4603+
ADDOP(c, elt_loc, POP_TOP);
4604+
}
4605+
break;
4606+
case COMP_LISTCOMP:
4607+
if (avoid_creation) {
4608+
if (elt->kind == Starred_kind) {
4609+
RETURN_IF_ERROR(codegen_unpack_starred(c, elt_loc, elt->v.Starred.value, /*yield=*/false));
4610+
} else {
4611+
VISIT(c, expr, elt);
4612+
ADDOP(c, elt_loc, POP_TOP);
4613+
}
4614+
break;
4615+
}
4616+
if (elt->kind == Starred_kind) {
4617+
VISIT(c, expr, elt->v.Starred.value);
4618+
ADDOP_I(c, elt_loc, LIST_EXTEND, depth + 1);
4619+
}
4620+
else {
4621+
VISIT(c, expr, elt);
4622+
ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
4623+
}
4624+
break;
4625+
case COMP_SETCOMP:
4626+
assert(!avoid_creation);
4627+
if (elt->kind == Starred_kind) {
4628+
VISIT(c, expr, elt->v.Starred.value);
4629+
ADDOP_I(c, elt_loc, SET_UPDATE, depth + 1);
4630+
}
4631+
else {
4632+
VISIT(c, expr, elt);
4633+
ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
4634+
}
4635+
break;
4636+
case COMP_DICTCOMP:
4637+
assert(!avoid_creation);
4638+
if (val == NULL) {
4639+
/* unpacking (**) case */
4640+
VISIT(c, expr, elt);
4641+
ADDOP_I(c, elt_loc, DICT_UPDATE, depth + 1);
4642+
}
4643+
else {
4644+
/* With '{k: v}', k is evaluated before v, so we do
4645+
the same. */
4646+
VISIT(c, expr, elt);
4647+
VISIT(c, expr, val);
4648+
elt_loc = LOCATION(elt->lineno,
4649+
val->end_lineno,
4650+
elt->col_offset,
4651+
val->end_col_offset);
4652+
ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
4653+
}
4654+
break;
4655+
default:
4656+
return ERROR;
4657+
}
4658+
return SUCCESS;
4659+
}
4660+
45894661
static int
45904662
codegen_sync_comprehension_generator(compiler *c, location loc,
45914663
asdl_comprehension_seq *generators,
@@ -4667,67 +4739,14 @@ codegen_sync_comprehension_generator(compiler *c, location loc,
46674739
/* only append after the last for generator */
46684740
if (gen_index >= asdl_seq_LEN(generators)) {
46694741
/* comprehension specific code */
4670-
switch (type) {
4671-
case COMP_GENEXP:
4672-
assert(!avoid_creation);
4673-
if (elt->kind == Starred_kind) {
4674-
RETURN_IF_ERROR(codegen_unpack_starred(c, elt_loc, elt->v.Starred.value, /*yield=*/true));
4675-
}
4676-
else {
4677-
VISIT(c, expr, elt);
4678-
ADDOP_YIELD(c, elt_loc);
4679-
ADDOP(c, elt_loc, POP_TOP);
4680-
}
4681-
break;
4682-
case COMP_LISTCOMP:
4683-
if (avoid_creation) {
4684-
if (elt->kind == Starred_kind) {
4685-
RETURN_IF_ERROR(codegen_unpack_starred(c, elt_loc, elt->v.Starred.value, /*yield=*/false));
4686-
} else {
4687-
VISIT(c, expr, elt);
4688-
ADDOP(c, elt_loc, POP_TOP);
4689-
}
4690-
break;
4691-
}
4692-
if (elt->kind == Starred_kind) {
4693-
VISIT(c, expr, elt->v.Starred.value);
4694-
ADDOP_I(c, elt_loc, LIST_EXTEND, depth + 1);
4695-
}
4696-
else {
4697-
VISIT(c, expr, elt);
4698-
ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
4699-
}
4700-
break;
4701-
case COMP_SETCOMP:
4702-
if (elt->kind == Starred_kind) {
4703-
VISIT(c, expr, elt->v.Starred.value);
4704-
ADDOP_I(c, elt_loc, SET_UPDATE, depth + 1);
4705-
}
4706-
else {
4707-
VISIT(c, expr, elt);
4708-
ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
4709-
}
4710-
break;
4711-
case COMP_DICTCOMP:
4712-
if (val == NULL) {
4713-
/* unpacking (**) case */
4714-
VISIT(c, expr, elt);
4715-
ADDOP_I(c, elt_loc, DICT_UPDATE, depth+1);
4716-
}
4717-
else {
4718-
/* With '{k: v}', k is evaluated before v, so we do
4719-
the same. */
4720-
VISIT(c, expr, elt);
4721-
VISIT(c, expr, val);
4722-
elt_loc = LOCATION(elt->lineno,
4723-
val->end_lineno,
4724-
elt->col_offset,
4725-
val->end_col_offset);
4726-
ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
4727-
}
4728-
break;
4729-
default:
4730-
return ERROR;
4742+
RETURN_IF_ERROR(codegen_comprehension_generator_helper(c, elt_loc, depth,
4743+
elt, val, type,
4744+
avoid_creation));
4745+
if (type == COMP_DICTCOMP && val != NULL) {
4746+
elt_loc = LOCATION(elt->lineno,
4747+
val->end_lineno,
4748+
elt->col_offset,
4749+
val->end_col_offset);
47314750
}
47324751
}
47334752

@@ -4809,81 +4828,14 @@ codegen_async_comprehension_generator(compiler *c, location loc,
48094828
/* only append after the last for generator */
48104829
if (gen_index >= asdl_seq_LEN(generators)) {
48114830
/* comprehension specific code */
4812-
switch (type) {
4813-
case COMP_GENEXP:
4814-
assert(!avoid_creation);
4815-
if (elt->kind == Starred_kind) {
4816-
NEW_JUMP_TARGET_LABEL(c, unpack_start);
4817-
NEW_JUMP_TARGET_LABEL(c, unpack_end);
4818-
VISIT(c, expr, elt->v.Starred.value);
4819-
ADDOP_I(c, elt_loc, GET_ITER, 0);
4820-
USE_LABEL(c, unpack_start);
4821-
ADDOP_JUMP(c, elt_loc, FOR_ITER, unpack_end);
4822-
ADDOP_YIELD(c, elt_loc);
4823-
ADDOP(c, elt_loc, POP_TOP);
4824-
ADDOP_JUMP(c, NO_LOCATION, JUMP, unpack_start);
4825-
USE_LABEL(c, unpack_end);
4826-
ADDOP(c, NO_LOCATION, END_FOR);
4827-
ADDOP(c, NO_LOCATION, POP_ITER);
4828-
}
4829-
else {
4830-
VISIT(c, expr, elt);
4831-
ADDOP_YIELD(c, elt_loc);
4832-
ADDOP(c, elt_loc, POP_TOP);
4833-
}
4834-
break;
4835-
case COMP_LISTCOMP:
4836-
if (avoid_creation) {
4837-
if (elt->kind == Starred_kind) {
4838-
RETURN_IF_ERROR(codegen_unpack_starred(c, elt_loc, elt->v.Starred.value, /*yield=*/false));
4839-
} else {
4840-
VISIT(c, expr, elt);
4841-
ADDOP(c, elt_loc, POP_TOP);
4842-
}
4843-
break;
4844-
}
4845-
4846-
if (elt->kind == Starred_kind) {
4847-
VISIT(c, expr, elt->v.Starred.value);
4848-
ADDOP_I(c, elt_loc, LIST_EXTEND, depth + 1);
4849-
}
4850-
else {
4851-
VISIT(c, expr, elt);
4852-
ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
4853-
}
4854-
break;
4855-
case COMP_SETCOMP:
4856-
assert(!avoid_creation);
4857-
if (elt->kind == Starred_kind) {
4858-
VISIT(c, expr, elt->v.Starred.value);
4859-
ADDOP_I(c, elt_loc, SET_UPDATE, depth + 1);
4860-
}
4861-
else {
4862-
VISIT(c, expr, elt);
4863-
ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
4864-
}
4865-
break;
4866-
case COMP_DICTCOMP:
4867-
assert(!avoid_creation);
4868-
if (val == NULL) {
4869-
/* unpacking (**) case */
4870-
VISIT(c, expr, elt);
4871-
ADDOP_I(c, elt_loc, DICT_UPDATE, depth+1);
4872-
}
4873-
else {
4874-
/* With '{k: v}', k is evaluated before v, so we do
4875-
the same. */
4876-
VISIT(c, expr, elt);
4877-
VISIT(c, expr, val);
4878-
elt_loc = LOCATION(elt->lineno,
4879-
val->end_lineno,
4880-
elt->col_offset,
4881-
val->end_col_offset);
4882-
ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
4883-
}
4884-
break;
4885-
default:
4886-
return ERROR;
4831+
RETURN_IF_ERROR(codegen_comprehension_generator_helper(c, elt_loc, depth,
4832+
elt, val, type,
4833+
avoid_creation));
4834+
if (type == COMP_DICTCOMP && val != NULL) {
4835+
elt_loc = LOCATION(elt->lineno,
4836+
val->end_lineno,
4837+
elt->col_offset,
4838+
val->end_col_offset);
48874839
}
48884840
}
48894841

0 commit comments

Comments
 (0)