Update the compiler package to compile the with-statement.

Jeremy, please review!

2 files changed, 62 insertions, 0 deletions

diff --git a/Lib/compiler/pycodegen.py b/Lib/compiler/pycodegen.py
index 4866e0e8cb..441ccf3c1b 100644
--- a/Lib/compiler/pycodegen.py
+++ b/Lib/compiler/pycodegen.py
@@ -807,6 +807,47 @@ class CodeGenerator:
         self.emit('END_FINALLY')
         self.setups.pop()
 
+    __with_count = 0
+
+    def visitWith(self, node):
+        body = self.newBlock()
+        final = self.newBlock()
+        exitvar = "$exit%d" % self.__with_count
+        valuevar = "$value%d" % self.__with_count
+        self.__with_count += 1
+        self.set_lineno(node)
+        self.visit(node.expr)
+        self.emit('LOAD_ATTR', '__context__')
+        self.emit('CALL_FUNCTION', 0)
+        self.emit('DUP_TOP')
+        self.emit('LOAD_ATTR', '__exit__')
+        self._implicitNameOp('STORE', exitvar)
+        self.emit('LOAD_ATTR', '__enter__')
+        self.emit('CALL_FUNCTION', 0)
+        if node.vars is None:
+            self.emit('POP_TOP')
+        else:
+            self._implicitNameOp('STORE', valuevar)
+        self.emit('SETUP_FINALLY', final)
+        self.nextBlock(body)
+        self.setups.push((TRY_FINALLY, body))
+        if node.vars is not None:
+            self._implicitNameOp('LOAD', valuevar)
+            self._implicitNameOp('DELETE', valuevar)
+            self.visit(node.vars)
+        self.visit(node.body)
+        self.emit('POP_BLOCK')
+        self.setups.pop()
+        self.emit('LOAD_CONST', None)
+        self.nextBlock(final)
+        self.setups.push((END_FINALLY, final))
+        self.emit('WITH_CLEANUP')
+        self.emit('CALL_FUNCTION', 3)
+        self.emit('POP_TOP')
+        self.emit('END_FINALLY')
+        self.setups.pop()
+        self.__with_count -= 1
+
     # misc
 
     def visitDiscard(self, node):
diff --git a/Lib/compiler/transformer.py b/Lib/compiler/transformer.py
index ae9b81903c..eed9ce93f7 100644
--- a/Lib/compiler/transformer.py
+++ b/Lib/compiler/transformer.py
@@ -536,6 +536,12 @@ class Transformer:
 
         return self.com_try_except(nodelist)
 
+    def with_stmt(self, nodelist):
+        return self.com_with(nodelist)
+
+    def with_var(self, nodelist):
+        return self.com_with_var(nodelist)
+
     def suite(self, nodelist):
         # simple_stmt | NEWLINE INDENT NEWLINE* (stmt NEWLINE*)+ DEDENT
         if len(nodelist) == 1:
@@ -926,6 +932,20 @@ class Transformer:
         return TryExcept(self.com_node(nodelist[2]), clauses, elseNode,
                          lineno=nodelist[0][2])
 
+    def com_with(self, nodelist):
+        # with_stmt: 'with' expr [with_var] ':' suite
+        expr = self.com_node(nodelist[1])
+        body = self.com_node(nodelist[-1])
+        if nodelist[2][0] == token.COLON:
+            var = None
+        else:
+            var = self.com_node(nodelist[2])
+        return With(expr, var, body, lineno=nodelist[0][2])
+
+    def com_with_var(self, nodelist):
+        # with_var: 'as' expr
+        return self.com_node(nodelist[1])
+
     def com_augassign_op(self, node):
         assert node[0] == symbol.augassign
         return node[1]
@@ -1390,6 +1410,7 @@ _legal_node_types = [
     symbol.while_stmt,
     symbol.for_stmt,
     symbol.try_stmt,
+    symbol.with_stmt,
     symbol.suite,
     symbol.testlist,
     symbol.testlist_safe,