Issue #2333: Backport set and dict comprehensions syntax.
diff --git a/Python/ast.c b/Python/ast.c
index 01d1f7a..c76014a 100644
--- a/Python/ast.c
+++ b/Python/ast.c
@@ -31,7 +31,7 @@
expr_context_ty);
static expr_ty ast_for_testlist(struct compiling *, const node *);
static stmt_ty ast_for_classdef(struct compiling *, const node *, asdl_seq *);
-static expr_ty ast_for_testlist_gexp(struct compiling *, const node *);
+static expr_ty ast_for_testlist_comp(struct compiling *, const node *);
/* Note different signature for ast_for_call */
static expr_ty ast_for_call(struct compiling *, const node *, expr_ty);
@@ -44,6 +44,9 @@
#define LINENO(n) ((n)->n_lineno)
#endif
+#define COMP_GENEXP 0
+#define COMP_SETCOMP 1
+
static identifier
new_identifier(const char* n, PyArena *arena) {
PyObject* id = PyString_InternFromString(n);
@@ -268,7 +271,7 @@
case eval_input: {
expr_ty testlist_ast;
- /* XXX Why not gen_for here? */
+ /* XXX Why not comp_for here? */
testlist_ast = ast_for_testlist(&c, CHILD(n, 0));
if (!testlist_ast)
goto error;
@@ -430,6 +433,12 @@
case ListComp_kind:
expr_name = "list comprehension";
break;
+ case SetComp_kind:
+ expr_name = "set comprehension";
+ break;
+ case DictComp_kind:
+ expr_name = "dict comprehension";
+ break;
case Dict_kind:
case Num_kind:
case Str_kind:
@@ -573,7 +582,7 @@
int i;
assert(TYPE(n) == testlist ||
TYPE(n) == listmaker ||
- TYPE(n) == testlist_gexp ||
+ TYPE(n) == testlist_comp ||
TYPE(n) == testlist_safe ||
TYPE(n) == testlist1);
@@ -1150,60 +1159,60 @@
return ListComp(elt, listcomps, LINENO(n), n->n_col_offset, c->c_arena);
}
-/* Count the number of 'for' loops in a generator expression.
+/*
+ Count the number of 'for' loops in a comprehension.
- Helper for ast_for_genexp().
+ Helper for ast_for_comprehension().
*/
static int
-count_gen_fors(struct compiling *c, const node *n)
+count_comp_fors(struct compiling *c, const node *n)
{
int n_fors = 0;
- node *ch = CHILD(n, 1);
- count_gen_for:
+ count_comp_for:
n_fors++;
- REQ(ch, gen_for);
- if (NCH(ch) == 5)
- ch = CHILD(ch, 4);
+ REQ(n, comp_for);
+ if (NCH(n) == 5)
+ n = CHILD(n, 4);
else
return n_fors;
- count_gen_iter:
- REQ(ch, gen_iter);
- ch = CHILD(ch, 0);
- if (TYPE(ch) == gen_for)
- goto count_gen_for;
- else if (TYPE(ch) == gen_if) {
- if (NCH(ch) == 3) {
- ch = CHILD(ch, 2);
- goto count_gen_iter;
+ count_comp_iter:
+ REQ(n, comp_iter);
+ n = CHILD(n, 0);
+ if (TYPE(n) == comp_for)
+ goto count_comp_for;
+ else if (TYPE(n) == comp_if) {
+ if (NCH(n) == 3) {
+ n = CHILD(n, 2);
+ goto count_comp_iter;
}
else
return n_fors;
}
-
+
/* Should never be reached */
PyErr_SetString(PyExc_SystemError,
- "logic error in count_gen_fors");
+ "logic error in count_comp_fors");
return -1;
}
-/* Count the number of 'if' statements in a generator expression.
+/* Count the number of 'if' statements in a comprehension.
- Helper for ast_for_genexp().
+ Helper for ast_for_comprehension().
*/
static int
-count_gen_ifs(struct compiling *c, const node *n)
+count_comp_ifs(struct compiling *c, const node *n)
{
int n_ifs = 0;
while (1) {
- REQ(n, gen_iter);
- if (TYPE(CHILD(n, 0)) == gen_for)
+ REQ(n, comp_iter);
+ if (TYPE(CHILD(n, 0)) == comp_for)
return n_ifs;
n = CHILD(n, 0);
- REQ(n, gen_if);
+ REQ(n, comp_if);
n_ifs++;
if (NCH(n) == 2)
return n_ifs;
@@ -1211,46 +1220,33 @@
}
}
-/* TODO(jhylton): Combine with list comprehension code? */
-static expr_ty
-ast_for_genexp(struct compiling *c, const node *n)
+static asdl_seq *
+ast_for_comprehension(struct compiling *c, const node *n)
{
- /* testlist_gexp: test ( gen_for | (',' test)* [','] )
- argument: [test '='] test [gen_for] # Really [keyword '='] test */
- expr_ty elt;
- asdl_seq *genexps;
int i, n_fors;
- node *ch;
-
- assert(TYPE(n) == (testlist_gexp) || TYPE(n) == (argument));
- assert(NCH(n) > 1);
-
- elt = ast_for_expr(c, CHILD(n, 0));
- if (!elt)
- return NULL;
-
- n_fors = count_gen_fors(c, n);
+ asdl_seq *comps;
+
+ n_fors = count_comp_fors(c, n);
if (n_fors == -1)
return NULL;
- genexps = asdl_seq_new(n_fors, c->c_arena);
- if (!genexps)
+ comps = asdl_seq_new(n_fors, c->c_arena);
+ if (!comps)
return NULL;
- ch = CHILD(n, 1);
for (i = 0; i < n_fors; i++) {
- comprehension_ty ge;
+ comprehension_ty comp;
asdl_seq *t;
expr_ty expression, first;
node *for_ch;
- REQ(ch, gen_for);
+ REQ(n, comp_for);
- for_ch = CHILD(ch, 1);
+ for_ch = CHILD(n, 1);
t = ast_for_exprlist(c, for_ch, Store);
if (!t)
return NULL;
- expression = ast_for_expr(c, CHILD(ch, 3));
+ expression = ast_for_expr(c, CHILD(n, 3));
if (!expression)
return NULL;
@@ -1258,21 +1254,20 @@
(x for x, in ...) has 1 element in t, but still requires a Tuple. */
first = (expr_ty)asdl_seq_GET(t, 0);
if (NCH(for_ch) == 1)
- ge = comprehension(first, expression, NULL, c->c_arena);
+ comp = comprehension(first, expression, NULL, c->c_arena);
else
- ge = comprehension(Tuple(t, Store, first->lineno, first->col_offset,
+ comp = comprehension(Tuple(t, Store, first->lineno, first->col_offset,
c->c_arena),
expression, NULL, c->c_arena);
-
- if (!ge)
+ if (!comp)
return NULL;
- if (NCH(ch) == 5) {
+ if (NCH(n) == 5) {
int j, n_ifs;
asdl_seq *ifs;
- ch = CHILD(ch, 4);
- n_ifs = count_gen_ifs(c, ch);
+ n = CHILD(n, 4);
+ n_ifs = count_comp_ifs(c, n);
if (n_ifs == -1)
return NULL;
@@ -1281,32 +1276,94 @@
return NULL;
for (j = 0; j < n_ifs; j++) {
- REQ(ch, gen_iter);
- ch = CHILD(ch, 0);
- REQ(ch, gen_if);
+ REQ(n, comp_iter);
+ n = CHILD(n, 0);
+ REQ(n, comp_if);
- expression = ast_for_expr(c, CHILD(ch, 1));
+ expression = ast_for_expr(c, CHILD(n, 1));
if (!expression)
return NULL;
asdl_seq_SET(ifs, j, expression);
- if (NCH(ch) == 3)
- ch = CHILD(ch, 2);
+ if (NCH(n) == 3)
+ n = CHILD(n, 2);
}
- /* on exit, must guarantee that ch is a gen_for */
- if (TYPE(ch) == gen_iter)
- ch = CHILD(ch, 0);
- ge->ifs = ifs;
+ /* on exit, must guarantee that n is a comp_for */
+ if (TYPE(n) == comp_iter)
+ n = CHILD(n, 0);
+ comp->ifs = ifs;
}
- asdl_seq_SET(genexps, i, ge);
+ asdl_seq_SET(comps, i, comp);
}
+ return comps;
+}
+
+static expr_ty
+ast_for_itercomp(struct compiling *c, const node *n, int type)
+{
+ expr_ty elt;
+ asdl_seq *comps;
- return GeneratorExp(elt, genexps, LINENO(n), n->n_col_offset, c->c_arena);
+ assert(NCH(n) > 1);
+
+ elt = ast_for_expr(c, CHILD(n, 0));
+ if (!elt)
+ return NULL;
+
+ comps = ast_for_comprehension(c, CHILD(n, 1));
+ if (!comps)
+ return NULL;
+
+ if (type == COMP_GENEXP)
+ return GeneratorExp(elt, comps, LINENO(n), n->n_col_offset, c->c_arena);
+ else if (type == COMP_SETCOMP)
+ return SetComp(elt, comps, LINENO(n), n->n_col_offset, c->c_arena);
+ else
+ /* Should never happen */
+ return NULL;
+}
+
+static expr_ty
+ast_for_dictcomp(struct compiling *c, const node *n)
+{
+ expr_ty key, value;
+ asdl_seq *comps;
+
+ assert(NCH(n) > 3);
+ REQ(CHILD(n, 1), COLON);
+
+ key = ast_for_expr(c, CHILD(n, 0));
+ if (!key)
+ return NULL;
+
+ value = ast_for_expr(c, CHILD(n, 2));
+ if (!value)
+ return NULL;
+
+ comps = ast_for_comprehension(c, CHILD(n, 3));
+ if (!comps)
+ return NULL;
+
+ return DictComp(key, value, comps, LINENO(n), n->n_col_offset, c->c_arena);
+}
+
+static expr_ty
+ast_for_genexp(struct compiling *c, const node *n)
+{
+ assert(TYPE(n) == (testlist_comp) || TYPE(n) == (argument));
+ return ast_for_itercomp(c, n, COMP_GENEXP);
+}
+
+static expr_ty
+ast_for_setcomp(struct compiling *c, const node *n)
+{
+ assert(TYPE(n) == (dictorsetmaker));
+ return ast_for_itercomp(c, n, COMP_SETCOMP);
}
static expr_ty
ast_for_atom(struct compiling *c, const node *n)
{
- /* atom: '(' [yield_expr|testlist_gexp] ')' | '[' [listmaker] ']'
+ /* atom: '(' [yield_expr|testlist_comp] ')' | '[' [listmaker] ']'
| '{' [dictmaker] '}' | '`' testlist '`' | NAME | NUMBER | STRING+
*/
node *ch = CHILD(n, 0);
@@ -1365,7 +1422,7 @@
if (TYPE(ch) == yield_expr)
return ast_for_expr(c, ch);
- return ast_for_testlist_gexp(c, ch);
+ return ast_for_testlist_comp(c, ch);
case LSQB: /* list (or list comprehension) */
ch = CHILD(n, 1);
@@ -1383,8 +1440,9 @@
else
return ast_for_listcomp(c, ch);
case LBRACE: {
- /* dictorsetmaker: test ':' test (',' test ':' test)* [','] |
- * test (',' test)* [','])
+ /* dictorsetmaker:
+ * (test ':' test (comp_for | (',' test ':' test)* [','])) |
+ * (test (comp_for | (',' test)* [',']))
*/
int i, size;
asdl_seq *keys, *values;
@@ -1408,6 +1466,11 @@
asdl_seq_SET(elts, i / 2, expression);
}
return Set(elts, LINENO(n), n->n_col_offset, c->c_arena);
+ } else if (TYPE(CHILD(ch, 1)) == comp_for) {
+ /* it's a set comprehension */
+ return ast_for_setcomp(c, ch);
+ } else if (NCH(ch) > 3 && TYPE(CHILD(ch, 3)) == comp_for) {
+ return ast_for_dictcomp(c, ch);
} else {
/* it's a dict */
size = (NCH(ch) + 1) / 4; /* +1 in case no trailing comma */
@@ -1916,7 +1979,7 @@
/*
arglist: (argument ',')* (argument [',']| '*' test [',' '**' test]
| '**' test)
- argument: [test '='] test [gen_for] # Really [keyword '='] test
+ argument: [test '='] test [comp_for] # Really [keyword '='] test
*/
int i, nargs, nkeywords, ngens;
@@ -1934,7 +1997,7 @@
if (TYPE(ch) == argument) {
if (NCH(ch) == 1)
nargs++;
- else if (TYPE(CHILD(ch, 1)) == gen_for)
+ else if (TYPE(CHILD(ch, 1)) == comp_for)
ngens++;
else
nkeywords++;
@@ -1979,7 +2042,7 @@
return NULL;
asdl_seq_SET(args, nargs++, e);
}
- else if (TYPE(CHILD(ch, 1)) == gen_for) {
+ else if (TYPE(CHILD(ch, 1)) == comp_for) {
e = ast_for_genexp(c, ch);
if (!e)
return NULL;
@@ -2049,14 +2112,14 @@
static expr_ty
ast_for_testlist(struct compiling *c, const node* n)
{
- /* testlist_gexp: test (',' test)* [','] */
+ /* testlist_comp: test (',' test)* [','] */
/* testlist: test (',' test)* [','] */
/* testlist_safe: test (',' test)+ [','] */
/* testlist1: test (',' test)* */
assert(NCH(n) > 0);
- if (TYPE(n) == testlist_gexp) {
+ if (TYPE(n) == testlist_comp) {
if (NCH(n) > 1)
- assert(TYPE(CHILD(n, 1)) != gen_for);
+ assert(TYPE(CHILD(n, 1)) != comp_for);
}
else {
assert(TYPE(n) == testlist ||
@@ -2074,12 +2137,12 @@
}
static expr_ty
-ast_for_testlist_gexp(struct compiling *c, const node* n)
+ast_for_testlist_comp(struct compiling *c, const node* n)
{
- /* testlist_gexp: test ( gen_for | (',' test)* [','] ) */
- /* argument: test [ gen_for ] */
- assert(TYPE(n) == testlist_gexp || TYPE(n) == argument);
- if (NCH(n) > 1 && TYPE(CHILD(n, 1)) == gen_for)
+ /* testlist_comp: test ( comp_for | (',' test)* [','] ) */
+ /* argument: test [ comp_for ] */
+ assert(TYPE(n) == testlist_comp || TYPE(n) == argument);
+ if (NCH(n) > 1 && TYPE(CHILD(n, 1)) == comp_for)
return ast_for_genexp(c, n);
return ast_for_testlist(c, n);
}