bpo-33041: Fixed bytecode generation for "async for" with a complex target. (#6052)

A StopAsyncIteration raised on assigning or unpacking will be now propagated
instead of stopping the iteration.

3 files changed, 70 insertions, 2 deletions

diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py
index b37b61b..f4a9d2a 100644
--- a/Lib/test/test_coroutines.py
+++ b/Lib/test/test_coroutines.py
@@ -1949,6 +1949,71 @@ class CoroutineTest(unittest.TestCase):
             support.gc_collect()
         self.assertIn("was never awaited", stderr.getvalue())
 
+    def test_for_assign_raising_stop_async_iteration(self):
+        class BadTarget:
+            def __setitem__(self, key, value):
+                raise StopAsyncIteration(42)
+        tgt = BadTarget()
+        async def source():
+            yield 10
+
+        async def run_for():
+            with self.assertRaises(StopAsyncIteration) as cm:
+                async for tgt[0] in source():
+                    pass
+            self.assertEqual(cm.exception.args, (42,))
+            return 'end'
+        self.assertEqual(run_async(run_for()), ([], 'end'))
+
+        async def run_list():
+            with self.assertRaises(StopAsyncIteration) as cm:
+                return [0 async for tgt[0] in source()]
+            self.assertEqual(cm.exception.args, (42,))
+            return 'end'
+        self.assertEqual(run_async(run_list()), ([], 'end'))
+
+        async def run_gen():
+            gen = (0 async for tgt[0] in source())
+            a = gen.asend(None)
+            with self.assertRaises(RuntimeError) as cm:
+                await a
+            self.assertIsInstance(cm.exception.__cause__, StopAsyncIteration)
+            self.assertEqual(cm.exception.__cause__.args, (42,))
+            return 'end'
+        self.assertEqual(run_async(run_gen()), ([], 'end'))
+
+    def test_for_assign_raising_stop_async_iteration_2(self):
+        class BadIterable:
+            def __iter__(self):
+                raise StopAsyncIteration(42)
+        async def badpairs():
+            yield BadIterable()
+
+        async def run_for():
+            with self.assertRaises(StopAsyncIteration) as cm:
+                async for i, j in badpairs():
+                    pass
+            self.assertEqual(cm.exception.args, (42,))
+            return 'end'
+        self.assertEqual(run_async(run_for()), ([], 'end'))
+
+        async def run_list():
+            with self.assertRaises(StopAsyncIteration) as cm:
+                return [0 async for i, j in badpairs()]
+            self.assertEqual(cm.exception.args, (42,))
+            return 'end'
+        self.assertEqual(run_async(run_list()), ([], 'end'))
+
+        async def run_gen():
+            gen = (0 async for i, j in badpairs())
+            a = gen.asend(None)
+            with self.assertRaises(RuntimeError) as cm:
+                await a
+            self.assertIsInstance(cm.exception.__cause__, StopAsyncIteration)
+            self.assertEqual(cm.exception.__cause__.args, (42,))
+            return 'end'
+        self.assertEqual(run_async(run_gen()), ([], 'end'))
+
 
 class CoroAsyncIOCompatTest(unittest.TestCase):
 
diff --git a/Misc/NEWS.d/next/Core and Builtins/2018-03-10-15-16-40.bpo-33041.-ak5Fk.rst b/Misc/NEWS.d/next/Core and Builtins/2018-03-10-15-16-40.bpo-33041.-ak5Fk.rst
new file mode 100644
index 0000000..af9ccfd
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2018-03-10-15-16-40.bpo-33041.-ak5Fk.rst
@@ -0,0 +1,3 @@
+Fixed bytecode generation for "async for" with a complex target. A
+StopAsyncIteration raised on assigning or unpacking will be now propagated
+instead of stopping the iteration.
diff --git a/Python/compile.c b/Python/compile.c
index f14647b..6c9e795 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -2473,8 +2473,8 @@ compiler_async_for(struct compiler *c, stmt_ty s)
     ADDOP(c, GET_ANEXT);
     ADDOP_O(c, LOAD_CONST, Py_None, consts);
     ADDOP(c, YIELD_FROM);
-    VISIT(c, expr, s->v.AsyncFor.target);
     ADDOP(c, POP_BLOCK);  /* for SETUP_FINALLY */
+    VISIT(c, expr, s->v.AsyncFor.target);
     compiler_pop_fblock(c, EXCEPT, try);
     ADDOP_JREL(c, JUMP_FORWARD, after_try);
 
@@ -4060,8 +4060,8 @@ compiler_async_comprehension_generator(struct compiler *c,
     ADDOP(c, GET_ANEXT);
     ADDOP_O(c, LOAD_CONST, Py_None, consts);
     ADDOP(c, YIELD_FROM);
-    VISIT(c, expr, gen->target);
     ADDOP(c, POP_BLOCK);
+    VISIT(c, expr, gen->target);
     compiler_pop_fblock(c, EXCEPT, try);
     ADDOP_JREL(c, JUMP_FORWARD, after_try);