Add the ability to parse set comprehensions

src/codegen/parse_ast.py

@@ -94,6 +94,8 @@ TYPE_MAP = {
         _ast.UAdd: 85,
         _ast.FloorDiv: 86,
         _ast.Ellipsis: 87,
+        _ast.Expression: 88,
+        _ast.SetComp: 89,
     }
 
 if sys.version_info >= (2,7):

src/codegen/parser.cpp

@@ -640,6 +640,16 @@ AST_Set* read_set(BufferedReader* reader) {
     return rtn;
 }
 
+AST_SetComp* read_setcomp(BufferedReader* reader) {
+    AST_SetComp* rtn = new AST_SetComp();
+
+    rtn->col_offset = readColOffset(reader);
+    rtn->elt = readASTExpr(reader);
+    readMiscVector(rtn->generators, reader);
+    rtn->lineno = reader->readULL();
+    return rtn;
+}
+
 AST_Slice* read_slice(BufferedReader* reader) {
     AST_Slice* rtn = new AST_Slice();
 
@@ -805,6 +815,8 @@ AST_expr* readASTExpr(BufferedReader* reader) {
             return read_repr(reader);
         case AST_TYPE::Set:
             return read_set(reader);
+        case AST_TYPE::SetComp:
+            return read_setcomp(reader);
         case AST_TYPE::Slice:
             return read_slice(reader);
         case AST_TYPE::Str:

src/codegen/pypa-parser.cpp

@@ -305,8 +305,7 @@ struct expr_dispatcher {
 
     template <typename T> ResultPtr read(T& item) {
         pypa::Ast& a = item;
-        fprintf(stderr, "Unhandled ast expression type caught: %d @%s\n", a.type, __PRETTY_FUNCTION__);
-        return nullptr;
+        RELEASE_ASSERT(0, "Unhandled ast expression type caught: %d @%s\n", a.type, __PRETTY_FUNCTION__);
     }
 
     ResultPtr read(pypa::AstAttribute& a) {
@@ -492,6 +491,14 @@ struct expr_dispatcher {
         return ptr;
     }
 
+    ResultPtr read(pypa::AstSetComp& l) {
+        AST_SetComp* ptr = new AST_SetComp();
+        location(ptr, l);
+        readVector(ptr->generators, l.generators, interned_strings);
+        ptr->elt = readItem(l.element, interned_strings);
+        return ptr;
+    }
+
     ResultPtr read(pypa::AstSlice& s) {
         AST_Slice* ptr = new AST_Slice();
         location(ptr, s);

src/core/ast.cpp

@@ -792,6 +792,22 @@ void* AST_Set::accept_expr(ExprVisitor* v) {
     return v->visit_set(this);
 }
 
+void AST_SetComp::accept(ASTVisitor* v) {
+    bool skip = v->visit_setcomp(this);
+    if (skip)
+        return;
+
+    for (auto c : generators) {
+        c->accept(v);
+    }
+
+    elt->accept(v);
+}
+
+void* AST_SetComp::accept_expr(ExprVisitor* v) {
+    return v->visit_setcomp(this);
+}
+
 void AST_Slice::accept(ASTVisitor* v) {
     bool skip = v->visit_slice(this);
     if (skip)
@@ -1610,6 +1626,17 @@ bool PrintVisitor::visit_set(AST_Set* node) {
     return true;
 }
 
+bool PrintVisitor::visit_setcomp(AST_SetComp* node) {
+    printf("{");
+    node->elt->accept(this);
+    for (auto c : node->generators) {
+        printf(" ");
+        c->accept(this);
+    }
+    printf("}");
+    return true;
+}
+
 bool PrintVisitor::visit_slice(AST_Slice* node) {
     printf("<slice>(");
     if (node->lower)
@@ -2011,6 +2038,10 @@ public:
         output->push_back(node);
         return false;
     }
+    virtual bool visit_setcomp(AST_SetComp* node) {
+        output->push_back(node);
+        return false;
+    }
     virtual bool visit_slice(AST_Slice* node) {
         output->push_back(node);
         return false;

src/core/ast.h

@@ -119,6 +119,7 @@ enum AST_TYPE {
     Set = 43,
     Ellipsis = 87,
     Expression = 88,
+    SetComp = 89,
 
     // Pseudo-nodes that are specific to this compiler:
     Branch = 200,
@@ -463,6 +464,21 @@ public:
     static const AST_TYPE::AST_TYPE TYPE = AST_TYPE::Exec;
 };
 
+// (Alternative to AST_Module, used for, e.g., eval)
+class AST_Expression : public AST {
+public:
+    std::unique_ptr<InternedStringPool> interned_strings;
+
+    AST_expr* body;
+
+    virtual void accept(ASTVisitor* v);
+
+    AST_Expression(std::unique_ptr<InternedStringPool> interned_strings)
+        : AST(AST_TYPE::Expression), interned_strings(std::move(interned_strings)) {}
+
+    static const AST_TYPE::AST_TYPE TYPE = AST_TYPE::Expression;
+};
+
 class AST_ExtSlice : public AST_expr {
 public:
     std::vector<AST_expr*> dims;
@@ -658,21 +674,6 @@ public:
     static const AST_TYPE::AST_TYPE TYPE = AST_TYPE::Module;
 };
 
-// (Alternative to AST_Module, used for, e.g., eval)
-class AST_Expression : public AST {
-public:
-    std::unique_ptr<InternedStringPool> interned_strings;
-
-    AST_expr* body;
-
-    virtual void accept(ASTVisitor* v);
-
-    AST_Expression(std::unique_ptr<InternedStringPool> interned_strings)
-        : AST(AST_TYPE::Expression), interned_strings(std::move(interned_strings)) {}
-
-    static const AST_TYPE::AST_TYPE TYPE = AST_TYPE::Expression;
-};
-
 class AST_Name : public AST_expr {
 public:
     AST_TYPE::AST_TYPE ctx_type;
@@ -801,6 +802,19 @@ public:
     static const AST_TYPE::AST_TYPE TYPE = AST_TYPE::Set;
 };
 
+class AST_SetComp : public AST_expr {
+public:
+    std::vector<AST_comprehension*> generators;
+    AST_expr* elt;
+
+    virtual void accept(ASTVisitor* v);
+    virtual void* accept_expr(ExprVisitor* v);
+
+    AST_SetComp() : AST_expr(AST_TYPE::SetComp) {}
+
+    const static AST_TYPE::AST_TYPE TYPE = AST_TYPE::SetComp;
+};
+
 class AST_Slice : public AST_expr {
 public:
     AST_expr* lower, *upper, *step;
@@ -1089,6 +1103,7 @@ public:
     virtual bool visit_repr(AST_Repr* node) { RELEASE_ASSERT(0, ""); }
     virtual bool visit_return(AST_Return* node) { RELEASE_ASSERT(0, ""); }
     virtual bool visit_set(AST_Set* node) { RELEASE_ASSERT(0, ""); }
+    virtual bool visit_setcomp(AST_SetComp* node) { RELEASE_ASSERT(0, ""); }
     virtual bool visit_slice(AST_Slice* node) { RELEASE_ASSERT(0, ""); }
     virtual bool visit_str(AST_Str* node) { RELEASE_ASSERT(0, ""); }
     virtual bool visit_subscript(AST_Subscript* node) { RELEASE_ASSERT(0, ""); }
@@ -1158,6 +1173,7 @@ public:
     virtual bool visit_repr(AST_Repr* node) { return false; }
     virtual bool visit_return(AST_Return* node) { return false; }
     virtual bool visit_set(AST_Set* node) { return false; }
+    virtual bool visit_setcomp(AST_SetComp* node) { return false; }
     virtual bool visit_slice(AST_Slice* node) { return false; }
     virtual bool visit_str(AST_Str* node) { return false; }
     virtual bool visit_subscript(AST_Subscript* node) { return false; }
@@ -1200,6 +1216,7 @@ public:
     virtual void* visit_num(AST_Num* node) { RELEASE_ASSERT(0, ""); }
     virtual void* visit_repr(AST_Repr* node) { RELEASE_ASSERT(0, ""); }
     virtual void* visit_set(AST_Set* node) { RELEASE_ASSERT(0, ""); }
+    virtual void* visit_setcomp(AST_SetComp* node) { RELEASE_ASSERT(0, ""); }
     virtual void* visit_slice(AST_Slice* node) { RELEASE_ASSERT(0, ""); }
     virtual void* visit_str(AST_Str* node) { RELEASE_ASSERT(0, ""); }
     virtual void* visit_subscript(AST_Subscript* node) { RELEASE_ASSERT(0, ""); }
@@ -1301,6 +1318,7 @@ public:
     virtual bool visit_repr(AST_Repr* node);
     virtual bool visit_return(AST_Return* node);
     virtual bool visit_set(AST_Set* node);
+    virtual bool visit_setcomp(AST_SetComp* node);
     virtual bool visit_slice(AST_Slice* node);
     virtual bool visit_str(AST_Str* node);
     virtual bool visit_subscript(AST_Subscript* node);

test/tests/set_comprehensions.py

+# For now, just check that we can parse set comprehensions even if we can't run them:
+
+def f():
+    print sorted({i%3 for i in xrange(5)})