Support for with statement.

ASTNode.cpp

@@ -79,7 +79,7 @@ const char* ASTBlock::type_str() const
 {
     static const char* s_type_strings[] = {
         "", "if", "else", "elif", "try", "CONTAINER", "except",
-        "finally", "while", "for"
+        "finally", "while", "for", "with",
     };
     return s_type_strings[blktype()];
 }

ASTNode.h

@@ -419,9 +419,9 @@ public:
     typedef std::list<PycRef<ASTNode> > list_t;
 
     enum BlkType {
-        BLK_MAIN, BLK_IF, BLK_ELSE, BLK_ELIF, BLK_TRY, 
+        BLK_MAIN, BLK_IF, BLK_ELSE, BLK_ELIF, BLK_TRY,
         BLK_CONTAINER, BLK_EXCEPT, BLK_FINALLY,
-        BLK_WHILE, BLK_FOR
+        BLK_WHILE, BLK_FOR, BLK_WITH
     };
 
     ASTBlock(BlkType blktype, int end = 0, int inited = 0)
@@ -505,6 +505,22 @@ private:
     int m_except;
 };
 
+class ASTWithBlock : public ASTBlock {
+public:
+    ASTWithBlock(int end)
+        : ASTBlock(ASTBlock::BLK_WITH, end) { }
+
+    PycRef<ASTNode> expr() const { return m_expr; }
+    PycRef<ASTNode> var() const { return m_var; }
+
+    void setExpr(PycRef<ASTNode> expr) { m_expr = expr; init(); }
+    void setVar(PycRef<ASTNode> var) { m_var = var; }
+
+private:
+    PycRef<ASTNode> m_expr;
+    PycRef<ASTNode> m_var;      // optional value
+};
+
 class ASTComprehension : public ASTNode {
 public:
     typedef std::list<PycRef<ASTIterBlock> > generator_t;

ASTree.cpp

@@ -70,8 +70,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                 && opcode != Pyc::POP_JUMP_IF_FALSE_A
                 && opcode != Pyc::JUMP_IF_TRUE_A
                 && opcode != Pyc::JUMP_IF_TRUE_OR_POP_A
-                && opcode != Pyc::POP_JUMP_IF_TRUE_A
-                && opcode != Pyc::POP_BLOCK) {
+                && opcode != Pyc::POP_JUMP_IF_TRUE_A) {
             else_pop = false;
 
             PycRef<ASTBlock> prev = curblock;
@@ -1336,7 +1335,11 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                 PycRef<ASTNode> value = stack.top();
                 stack.pop();
                 if (!curblock->inited()) {
-                    curblock.cast<ASTCondBlock>()->init();
+                    if (curblock->blktype() == ASTBlock::BLK_WITH) {
+                        curblock.cast<ASTWithBlock>()->setExpr(value);
+                    } else {
+                        curblock.cast<ASTCondBlock>()->init();
+                    }
                     break;
                 } else if (value->type() == ASTNode::NODE_INVALID
                         || value->type() == ASTNode::NODE_BINARY
@@ -1473,6 +1476,24 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
         case Pyc::SET_LINENO_A:
             // Ignore
             break;
+        case Pyc::SETUP_WITH_A:
+            {
+                PycRef<ASTBlock> withblock = new ASTWithBlock(pos+operand);
+                blocks.push(withblock);
+                curblock = blocks.top();
+            }
+            break;
+        case Pyc::WITH_CLEANUP:
+            {
+                // Stack top should be a None.
+                PycRef<ASTNode> none = stack.top();
+                stack.pop();
+
+                if (none != Node_NULL) {
+                    fprintf(stderr, "Something TERRIBLE happened!\n");
+                }
+            }
+            break;
         case Pyc::SETUP_EXCEPT_A:
             {
                 if (curblock->blktype() == ASTBlock::BLK_CONTAINER) {
@@ -1671,6 +1692,9 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                     if (curblock->blktype() == ASTBlock::BLK_FOR
                             && !curblock->inited()) {
                         curblock.cast<ASTIterBlock>()->setIndex(name);
+                    } else if (curblock->blktype() == ASTBlock::BLK_WITH) {
+                        curblock.cast<ASTWithBlock>()->setExpr(value);
+                        curblock.cast<ASTWithBlock>()->setVar(name);
                     } else {
                         curblock->append(new ASTStore(value, name));
                     }
@@ -1766,6 +1790,10 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                         PycRef<ASTImport> import = stack.top().cast<ASTImport>();
 
                         import->add_store(new ASTStore(value, name));
+                    } else if (curblock->blktype() == ASTBlock::BLK_WITH
+                               && !curblock->inited()) {
+                        curblock.cast<ASTWithBlock>()->setExpr(value);
+                        curblock.cast<ASTWithBlock>()->setVar(name);
                     } else {
                         curblock->append(new ASTStore(value, name));
 
@@ -2273,6 +2301,14 @@ void print_src(PycRef<ASTNode> node, PycModule* mod)
                     blk.cast<ASTCondBlock>()->cond() != Node_NULL) {
                 fprintf(pyc_output, " ");
                 print_src(blk.cast<ASTCondBlock>()->cond(), mod);
+            } else if (blk->blktype() == ASTBlock::BLK_WITH) {
+                fprintf(pyc_output, " ");
+                print_src(blk.cast<ASTWithBlock>()->expr(), mod);
+                PycRef<ASTNode> var = blk.cast<ASTWithBlock>()->var();
+                if (var != Node_NULL) {
+                    fprintf(pyc_output, " as ");
+                    print_src(var, mod);
+                }
             }
             fprintf(pyc_output, ":\n");
 

tests/test_dict.py

@@ -0,0 +1,4 @@
+x = { 'one': 1, 'two': 2 }
+if x['one'] == 1:
+    print("Correct!")
+x['three'] = 3

tests/test_with.py

@@ -0,0 +1,9 @@
+with open(__file__) as f:
+    f.read()
+
+with open(__file__):
+    x = 1
+    if x == 1:
+        print('one')
+    else:
+        x = 2