Merge pull request #65 from vinzenz/dict-comprehension

Implementation of dict comprehension

src/analysis/scoping_analysis.cpp

@@ -158,6 +158,7 @@ public:
     // virtual bool visit_classdef(AST_ClassDef *node) { return false; }
     virtual bool visit_continue(AST_Continue* node) { return false; }
     virtual bool visit_dict(AST_Dict* node) { return false; }
+    virtual bool visit_dictcomp(AST_DictComp* node) { return false; }
     virtual bool visit_excepthandler(AST_ExceptHandler* node) { return false; }
     virtual bool visit_expr(AST_Expr* node) { return false; }
     virtual bool visit_for(AST_For* node) { return false; }

src/codegen/parser.cpp

@@ -328,6 +328,17 @@ AST_Dict* read_dict(BufferedReader* reader) {
     return rtn;
 }
 
+AST_DictComp* read_dictcomp(BufferedReader* reader) {
+    AST_DictComp* rtn = new AST_DictComp();
+    rtn->col_offset = readColOffset(reader);
+    readMiscVector(rtn->generators, reader);
+    rtn->key = readASTExpr(reader);
+    rtn->lineno = reader->readULL();
+    rtn->value = readASTExpr(reader);
+    return rtn;
+}
+
+
 AST_ExceptHandler* read_excepthandler(BufferedReader* reader) {
     AST_ExceptHandler* rtn = new AST_ExceptHandler();
 
@@ -679,6 +690,8 @@ AST_expr* readASTExpr(BufferedReader* reader) {
             return read_compare(reader);
         case AST_TYPE::Dict:
             return read_dict(reader);
+        case AST_TYPE::DictComp:
+            return read_dictcomp(reader);
         case AST_TYPE::IfExp:
             return read_ifexp(reader);
         case AST_TYPE::Index:

src/core/ast.cpp

@@ -405,6 +405,23 @@ void* AST_Dict::accept_expr(ExprVisitor* v) {
     return v->visit_dict(this);
 }
 
+void AST_DictComp::accept(ASTVisitor* v) {
+    bool skip = v->visit_dictcomp(this);
+    if (skip)
+        return;
+
+    for (auto c : generators) {
+        c->accept(v);
+    }
+
+    value->accept(v);
+    key->accept(v);
+}
+
+void* AST_DictComp::accept_expr(ExprVisitor* v) {
+    return v->visit_dictcomp(this);
+}
+
 void AST_ExceptHandler::accept(ASTVisitor* v) {
     bool skip = v->visit_excepthandler(this);
     if (skip)
@@ -1100,6 +1117,19 @@ bool PrintVisitor::visit_dict(AST_Dict* node) {
     return true;
 }
 
+bool PrintVisitor::visit_dictcomp(AST_DictComp* node) {
+    printf("{");
+    node->key->accept(this);
+    printf(":");
+    node->value->accept(this);
+    for (auto c : node->generators) {
+        printf(" ");
+        c->accept(this);
+    }
+    printf("}");
+    return true;
+}
+
 bool PrintVisitor::visit_excepthandler(AST_ExceptHandler* node) {
     printf("except");
     if (node->type) {
@@ -1633,6 +1663,10 @@ public:
         output->push_back(node);
         return false;
     }
+    virtual bool visit_dictcomp(AST_DictComp* node) {
+        output->push_back(node);
+        return false;
+    }
     virtual bool visit_excepthandler(AST_ExceptHandler* node) {
         output->push_back(node);
         return false;

src/core/ast.h

@@ -364,6 +364,19 @@ public:
     static const AST_TYPE::AST_TYPE TYPE = AST_TYPE::Dict;
 };
 
+class AST_DictComp : public AST_expr {
+public:
+    std::vector<AST_comprehension*> generators;
+    AST_expr* key, *value;
+
+    virtual void accept(ASTVisitor* v);
+    virtual void* accept_expr(ExprVisitor* v);
+
+    AST_DictComp() : AST_expr(AST_TYPE::DictComp) {}
+
+    const static AST_TYPE::AST_TYPE TYPE = AST_TYPE::DictComp;
+};
+
 class AST_Delete : public AST_stmt {
 public:
     std::vector<AST_expr*> targets;
@@ -875,6 +888,7 @@ public:
     virtual bool visit_continue(AST_Continue* node) { RELEASE_ASSERT(0, ""); }
     virtual bool visit_delete(AST_Delete* node) { RELEASE_ASSERT(0, ""); }
     virtual bool visit_dict(AST_Dict* node) { RELEASE_ASSERT(0, ""); }
+    virtual bool visit_dictcomp(AST_DictComp* node) { RELEASE_ASSERT(0, ""); }
     virtual bool visit_excepthandler(AST_ExceptHandler* node) { RELEASE_ASSERT(0, ""); }
     virtual bool visit_expr(AST_Expr* node) { RELEASE_ASSERT(0, ""); }
     virtual bool visit_for(AST_For* node) { RELEASE_ASSERT(0, ""); }
@@ -935,6 +949,7 @@ public:
     virtual bool visit_continue(AST_Continue* node) { return false; }
     virtual bool visit_delete(AST_Delete* node) { return false; }
     virtual bool visit_dict(AST_Dict* node) { return false; }
+    virtual bool visit_dictcomp(AST_DictComp* node) { return false; }
     virtual bool visit_excepthandler(AST_ExceptHandler* node) { return false; }
     virtual bool visit_expr(AST_Expr* node) { return false; }
     virtual bool visit_for(AST_For* node) { return false; }
@@ -985,6 +1000,7 @@ public:
     virtual void* visit_clsattribute(AST_ClsAttribute* node) { RELEASE_ASSERT(0, ""); }
     virtual void* visit_compare(AST_Compare* node) { RELEASE_ASSERT(0, ""); }
     virtual void* visit_dict(AST_Dict* node) { RELEASE_ASSERT(0, ""); }
+    virtual void* visit_dictcomp(AST_DictComp* node) { RELEASE_ASSERT(0, ""); }
     virtual void* visit_ifexp(AST_IfExp* node) { RELEASE_ASSERT(0, ""); }
     virtual void* visit_index(AST_Index* node) { RELEASE_ASSERT(0, ""); }
     virtual void* visit_langprimitive(AST_LangPrimitive* node) { RELEASE_ASSERT(0, ""); }
@@ -1061,6 +1077,7 @@ public:
     virtual bool visit_continue(AST_Continue* node);
     virtual bool visit_delete(AST_Delete* node);
     virtual bool visit_dict(AST_Dict* node);
+    virtual bool visit_dictcomp(AST_DictComp* node);
     virtual bool visit_excepthandler(AST_ExceptHandler* node);
     virtual bool visit_expr(AST_Expr* node);
     virtual bool visit_for(AST_For* node);

test/tests/dictcomp.py

+def dict2str(d):
+    result = ''
+    for k, v in sorted(d.items()):
+        if result:
+            result += ', '
+        result += '%s: %s' % (str(k), str(v))
+    return '{%s}' % result
+
+print dict2str({i: j for i in range(4) for j in range(4)})
+
+def f():
+    print dict2str({i: j for i in range(4) for j in range(4)})
+    # print i, j
+f()
+
+
+# Combine a list comprehension with a bunch of other control-flow expressions:
+def f(x, y):
+    # TODO make sure to use an 'if' in a comprehension where the if contains control flow
+    print dict2str({y if i % 3 else y ** 2 + i: (i if i%2 else i/2) for i in (xrange(4 if x else 5) if y else xrange(3))})
+f(0, 0)
+f(0, 1)
+f(1, 0)
+f(1, 1)
+
+# TODO: test on ifs
+
+def f():
+    print dict2str({i : j for (i, j) in sorted({1:2, 3:4, 5:6, 7:8}.items())})
+f()
+
+# The expr should not get evaluated if the if-condition fails:
+def f():
+    def p(i):
+        print i
+        return i ** 2
+    def k(i):
+        print i
+        return i * 4 + i
+
+    print dict2str({k(i):p(i) for i in xrange(50) if i % 5 == 0 if i % 3 == 0})
+f()
+
+def f():
+    print dict2str({i: j for i in xrange(4) for j in xrange(i)})
+f()
+
+def f():
+    j = 1
+    # The 'if' part of this list comprehension references j;
+    # the first time through it will use the j above, but later times
+    # it may-or-may-not use the j from the inner part of the listcomp.
+    print dict2str({i: j for i in xrange(7) if i % 2 != j % 2 for j in xrange(i)})
+    f()
+
+def f():
+    # Checking the order of evaluation of the if conditions:
+
+    def c1(x):
+        print "c1", x
+        return x % 2 == 0
+
+    def c2(x):
+        print "c2", x
+        return x % 3 == 0
+
+    print dict2str({i : i for i in xrange(20) if c1(i) if c2(i)})
+f()
+
+def control_flow_in_listcomp():
+    print dict2str({(i ** 2 if i > 5 else i ** 2 * -1):(i if i else -1) for i in (xrange(10) if True else []) if (i % 2 == 0 or i % 3 != 0)})
+control_flow_in_listcomp()