Skip to content

Commit 109a192

Browse files
authored
gh-151907: Reduce duplication in codegen comprehension functions (GH-154846)
1 parent 9def455 commit 109a192

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
@@ -4625,6 +4625,78 @@ codegen_unpack_starred(compiler *c, location loc, expr_ty value, bool yield)
46254625
return SUCCESS;
46264626
}
46274627

4628+
static int
4629+
codegen_comprehension_generator_helper(compiler *c, location elt_loc, int depth,
4630+
expr_ty elt, expr_ty val, int type,
4631+
bool avoid_creation)
4632+
{
4633+
switch (type) {
4634+
case COMP_GENEXP:
4635+
assert(!avoid_creation);
4636+
if (elt->kind == Starred_kind) {
4637+
RETURN_IF_ERROR(codegen_unpack_starred(c, elt_loc, elt->v.Starred.value, /*yield=*/true));
4638+
}
4639+
else {
4640+
VISIT(c, expr, elt);
4641+
ADDOP_YIELD(c, elt_loc);
4642+
ADDOP(c, elt_loc, POP_TOP);
4643+
}
4644+
break;
4645+
case COMP_LISTCOMP:
4646+
if (avoid_creation) {
4647+
if (elt->kind == Starred_kind) {
4648+
RETURN_IF_ERROR(codegen_unpack_starred(c, elt_loc, elt->v.Starred.value, /*yield=*/false));
4649+
} else {
4650+
VISIT(c, expr, elt);
4651+
ADDOP(c, elt_loc, POP_TOP);
4652+
}
4653+
break;
4654+
}
4655+
if (elt->kind == Starred_kind) {
4656+
VISIT(c, expr, elt->v.Starred.value);
4657+
ADDOP_I(c, elt_loc, LIST_EXTEND, depth + 1);
4658+
}
4659+
else {
4660+
VISIT(c, expr, elt);
4661+
ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
4662+
}
4663+
break;
4664+
case COMP_SETCOMP:
4665+
assert(!avoid_creation);
4666+
if (elt->kind == Starred_kind) {
4667+
VISIT(c, expr, elt->v.Starred.value);
4668+
ADDOP_I(c, elt_loc, SET_UPDATE, depth + 1);
4669+
}
4670+
else {
4671+
VISIT(c, expr, elt);
4672+
ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
4673+
}
4674+
break;
4675+
case COMP_DICTCOMP:
4676+
assert(!avoid_creation);
4677+
if (val == NULL) {
4678+
/* unpacking (**) case */
4679+
VISIT(c, expr, elt);
4680+
ADDOP_I(c, elt_loc, DICT_UPDATE, depth + 1);
4681+
}
4682+
else {
4683+
/* With '{k: v}', k is evaluated before v, so we do
4684+
the same. */
4685+
VISIT(c, expr, elt);
4686+
VISIT(c, expr, val);
4687+
elt_loc = LOCATION(elt->lineno,
4688+
val->end_lineno,
4689+
elt->col_offset,
4690+
val->end_col_offset);
4691+
ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
4692+
}
4693+
break;
4694+
default:
4695+
return ERROR;
4696+
}
4697+
return SUCCESS;
4698+
}
4699+
46284700
static int
46294701
codegen_sync_comprehension_generator(compiler *c, location loc,
46304702
asdl_comprehension_seq *generators,
@@ -4706,67 +4778,14 @@ codegen_sync_comprehension_generator(compiler *c, location loc,
47064778
/* only append after the last for generator */
47074779
if (gen_index >= asdl_seq_LEN(generators)) {
47084780
/* comprehension specific code */
4709-
switch (type) {
4710-
case COMP_GENEXP:
4711-
assert(!avoid_creation);
4712-
if (elt->kind == Starred_kind) {
4713-
RETURN_IF_ERROR(codegen_unpack_starred(c, elt_loc, elt->v.Starred.value, /*yield=*/true));
4714-
}
4715-
else {
4716-
VISIT(c, expr, elt);
4717-
ADDOP_YIELD(c, elt_loc);
4718-
ADDOP(c, elt_loc, POP_TOP);
4719-
}
4720-
break;
4721-
case COMP_LISTCOMP:
4722-
if (avoid_creation) {
4723-
if (elt->kind == Starred_kind) {
4724-
RETURN_IF_ERROR(codegen_unpack_starred(c, elt_loc, elt->v.Starred.value, /*yield=*/false));
4725-
} else {
4726-
VISIT(c, expr, elt);
4727-
ADDOP(c, elt_loc, POP_TOP);
4728-
}
4729-
break;
4730-
}
4731-
if (elt->kind == Starred_kind) {
4732-
VISIT(c, expr, elt->v.Starred.value);
4733-
ADDOP_I(c, elt_loc, LIST_EXTEND, depth + 1);
4734-
}
4735-
else {
4736-
VISIT(c, expr, elt);
4737-
ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
4738-
}
4739-
break;
4740-
case COMP_SETCOMP:
4741-
if (elt->kind == Starred_kind) {
4742-
VISIT(c, expr, elt->v.Starred.value);
4743-
ADDOP_I(c, elt_loc, SET_UPDATE, depth + 1);
4744-
}
4745-
else {
4746-
VISIT(c, expr, elt);
4747-
ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
4748-
}
4749-
break;
4750-
case COMP_DICTCOMP:
4751-
if (val == NULL) {
4752-
/* unpacking (**) case */
4753-
VISIT(c, expr, elt);
4754-
ADDOP_I(c, elt_loc, DICT_UPDATE, depth+1);
4755-
}
4756-
else {
4757-
/* With '{k: v}', k is evaluated before v, so we do
4758-
the same. */
4759-
VISIT(c, expr, elt);
4760-
VISIT(c, expr, val);
4761-
elt_loc = LOCATION(elt->lineno,
4762-
val->end_lineno,
4763-
elt->col_offset,
4764-
val->end_col_offset);
4765-
ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
4766-
}
4767-
break;
4768-
default:
4769-
return ERROR;
4781+
RETURN_IF_ERROR(codegen_comprehension_generator_helper(c, elt_loc, depth,
4782+
elt, val, type,
4783+
avoid_creation));
4784+
if (type == COMP_DICTCOMP && val != NULL) {
4785+
elt_loc = LOCATION(elt->lineno,
4786+
val->end_lineno,
4787+
elt->col_offset,
4788+
val->end_col_offset);
47704789
}
47714790
}
47724791

@@ -4848,81 +4867,14 @@ codegen_async_comprehension_generator(compiler *c, location loc,
48484867
/* only append after the last for generator */
48494868
if (gen_index >= asdl_seq_LEN(generators)) {
48504869
/* comprehension specific code */
4851-
switch (type) {
4852-
case COMP_GENEXP:
4853-
assert(!avoid_creation);
4854-
if (elt->kind == Starred_kind) {
4855-
NEW_JUMP_TARGET_LABEL(c, unpack_start);
4856-
NEW_JUMP_TARGET_LABEL(c, unpack_end);
4857-
VISIT(c, expr, elt->v.Starred.value);
4858-
ADDOP_I(c, elt_loc, GET_ITER, 0);
4859-
USE_LABEL(c, unpack_start);
4860-
ADDOP_JUMP(c, elt_loc, FOR_ITER, unpack_end);
4861-
ADDOP_YIELD(c, elt_loc);
4862-
ADDOP(c, elt_loc, POP_TOP);
4863-
ADDOP_JUMP(c, NO_LOCATION, JUMP, unpack_start);
4864-
USE_LABEL(c, unpack_end);
4865-
ADDOP(c, NO_LOCATION, END_FOR);
4866-
ADDOP(c, NO_LOCATION, POP_ITER);
4867-
}
4868-
else {
4869-
VISIT(c, expr, elt);
4870-
ADDOP_YIELD(c, elt_loc);
4871-
ADDOP(c, elt_loc, POP_TOP);
4872-
}
4873-
break;
4874-
case COMP_LISTCOMP:
4875-
if (avoid_creation) {
4876-
if (elt->kind == Starred_kind) {
4877-
RETURN_IF_ERROR(codegen_unpack_starred(c, elt_loc, elt->v.Starred.value, /*yield=*/false));
4878-
} else {
4879-
VISIT(c, expr, elt);
4880-
ADDOP(c, elt_loc, POP_TOP);
4881-
}
4882-
break;
4883-
}
4884-
4885-
if (elt->kind == Starred_kind) {
4886-
VISIT(c, expr, elt->v.Starred.value);
4887-
ADDOP_I(c, elt_loc, LIST_EXTEND, depth + 1);
4888-
}
4889-
else {
4890-
VISIT(c, expr, elt);
4891-
ADDOP_I(c, elt_loc, LIST_APPEND, depth + 1);
4892-
}
4893-
break;
4894-
case COMP_SETCOMP:
4895-
assert(!avoid_creation);
4896-
if (elt->kind == Starred_kind) {
4897-
VISIT(c, expr, elt->v.Starred.value);
4898-
ADDOP_I(c, elt_loc, SET_UPDATE, depth + 1);
4899-
}
4900-
else {
4901-
VISIT(c, expr, elt);
4902-
ADDOP_I(c, elt_loc, SET_ADD, depth + 1);
4903-
}
4904-
break;
4905-
case COMP_DICTCOMP:
4906-
assert(!avoid_creation);
4907-
if (val == NULL) {
4908-
/* unpacking (**) case */
4909-
VISIT(c, expr, elt);
4910-
ADDOP_I(c, elt_loc, DICT_UPDATE, depth+1);
4911-
}
4912-
else {
4913-
/* With '{k: v}', k is evaluated before v, so we do
4914-
the same. */
4915-
VISIT(c, expr, elt);
4916-
VISIT(c, expr, val);
4917-
elt_loc = LOCATION(elt->lineno,
4918-
val->end_lineno,
4919-
elt->col_offset,
4920-
val->end_col_offset);
4921-
ADDOP_I(c, elt_loc, MAP_ADD, depth + 1);
4922-
}
4923-
break;
4924-
default:
4925-
return ERROR;
4870+
RETURN_IF_ERROR(codegen_comprehension_generator_helper(c, elt_loc, depth,
4871+
elt, val, type,
4872+
avoid_creation));
4873+
if (type == COMP_DICTCOMP && val != NULL) {
4874+
elt_loc = LOCATION(elt->lineno,
4875+
val->end_lineno,
4876+
elt->col_offset,
4877+
val->end_col_offset);
49264878
}
49274879
}
49284880

0 commit comments

Comments
 (0)