Issue #2333: Backport set and dict comprehensions syntax.

1 files changed, 149 insertions, 86 deletions

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 @@ static asdl_seq *ast_for_exprlist(struct compiling *, const node *,
                                   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 @@ static PyObject *parsestrplus(struct compiling *, const node *n);
 #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 @@ PyAST_FromNode(const node *n, PyCompilerFlags *flags, const char *filename,
         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 @@ set_context(struct compiling *c, expr_ty e, expr_context_ty ctx, const node *n)
         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 @@ seq_for_testlist(struct compiling *c, const node *n)
     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 @@ ast_for_listcomp(struct compiling *c, const node *n)
     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 @@ count_gen_ifs(struct compiling *c, const node *n)
     }
 }
 
-/* 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 @@ ast_for_genexp(struct compiling *c, const node *n)
            (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 @@ ast_for_genexp(struct compiling *c, const node *n)
                 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 @@ ast_for_atom(struct compiling *c, const node *n)
         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 @@ ast_for_atom(struct compiling *c, const node *n)
         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 @@ ast_for_atom(struct compiling *c, const node *n)
                 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 @@ ast_for_call(struct compiling *c, const node *n, expr_ty func)
     /*
       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 @@ ast_for_call(struct compiling *c, const node *n, expr_ty func)
         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 @@ ast_for_call(struct compiling *c, const node *n, expr_ty func)
                     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 @@ ast_for_call(struct compiling *c, const node *n, expr_ty func)
 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 @@ ast_for_testlist(struct compiling *c, const node* n)
 }
 
 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);
 }