[pypy-commit] pypy remove-del-from-generatoriterator: handle yield inside with and also nested blocks

[jerith]

Author: Jeremy Thurgood <firxen at gmail.com>
Branch: remove-del-from-generatoriterator
Changeset: r68720:e25d33c7d99a
Date: 2014-01-17 16:28 +0200
http://bitbucket.org/pypy/pypy/changeset/e25d33c7d99a/

Log:	handle yield inside with and also nested blocks

diff --git a/pypy/interpreter/astcompiler/symtable.py b/pypy/interpreter/astcompiler/symtable.py
--- a/pypy/interpreter/astcompiler/symtable.py
+++ b/pypy/interpreter/astcompiler/symtable.py
@@ -43,7 +43,7 @@
         self.child_has_free = False
         self.nested = False
         self.doc_removable = False
-        self._visiting_try_body = False
+        self._in_try_body_depth = 0
 
     def lookup(self, name):
         """Find the scope of identifier 'name'."""
@@ -78,11 +78,11 @@
 
     def note_try_start(self, try_node):
         """Called when a try is found, before visiting the body."""
-        self._visiting_try_body = True
+        self._in_try_body_depth += 1
 
     def note_try_end(self, try_node):
         """Called after visiting a try body."""
-        self._visiting_try_body = False
+        self._in_try_body_depth -= 1
 
     def note_yield(self, yield_node):
         """Called when a yield is found."""
@@ -230,7 +230,7 @@
             raise SyntaxError("'return' with argument inside generator",
                               self.ret.lineno, self.ret.col_offset)
         self.is_generator = True
-        if self._visiting_try_body:
+        if self._in_try_body_depth > 0:
             self.has_yield_inside_try = True
 
     def note_return(self, ret):
@@ -475,7 +475,12 @@
         self.scope.new_temporary_name()
         if wih.optional_vars:
             self.scope.new_temporary_name()
-        ast.GenericASTVisitor.visit_With(self, wih)
+        wih.context_expr.walkabout(self)
+        if wih.optional_vars:
+            wih.optional_vars.walkabout(self)
+        self.scope.note_try_start(wih)
+        self.visit_sequence(wih.body)
+        self.scope.note_try_end(wih)
 
     def visit_arguments(self, arguments):
         scope = self.scope
diff --git a/pypy/interpreter/astcompiler/test/test_symtable.py b/pypy/interpreter/astcompiler/test/test_symtable.py
--- a/pypy/interpreter/astcompiler/test/test_symtable.py
+++ b/pypy/interpreter/astcompiler/test/test_symtable.py
@@ -353,6 +353,15 @@
         assert scp.has_yield_inside_try
         scp = self.func_scope("def f():\n  try:\n    yield x\n  finally: pass")
         assert scp.has_yield_inside_try
+        scp = self.func_scope("def f():\n    with x: yield y")
+        assert scp.has_yield_inside_try
+
+    def test_yield_outside_try(self):
+        for input in ("try: pass\n    except: pass",
+                      "try: pass\n    finally: pass",
+                      "with x: pass"):
+            input = "def f():\n    yield y\n    %s\n    yield y" % (input,)
+            assert not self.func_scope(input).has_yield_inside_try
 
     def test_return(self):
         for input in ("class x: return", "return"):