diff --git a/Lib/compiler/pycodegen.py b/Lib/compiler/pycodegen.py index 4866e0e8cb6..441ccf3c1be 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 ae9b81903cc..eed9ce93f7a 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, diff --git a/Lib/test/test_compiler.py b/Lib/test/test_compiler.py index d2f062c928d..c328d714d03 100644 --- a/Lib/test/test_compiler.py +++ b/Lib/test/test_compiler.py @@ -20,7 +20,7 @@ class CompilerTest(unittest.TestCase): for basename in os.listdir(dir): if not basename.endswith(".py"): continue - if not TEST_ALL and random() < 0.98: + if not TEST_ALL and random() < 0.98 and basename != "test_with.py": continue path = os.path.join(dir, basename) if test.test_support.verbose: