bpo-33041: Fixed jumping if the function contains an "async for" loop. (GH-6154)

Lib/test/test_coroutines.py

@@ -1846,6 +1846,36 @@ class CoroutineTest(unittest.TestCase):
             run_async(run_gen()),
             ([], [121]))
 
+    def test_comp_4_2(self):
+        async def f(it):
+            for i in it:
+                yield i
+
+        async def run_list():
+            return [i + 10 async for i in f(range(5)) if 0 < i < 4]
+        self.assertEqual(
+            run_async(run_list()),
+            ([], [11, 12, 13]))
+
+        async def run_set():
+            return {i + 10 async for i in f(range(5)) if 0 < i < 4}
+        self.assertEqual(
+            run_async(run_set()),
+            ([], {11, 12, 13}))
+
+        async def run_dict():
+            return {i + 10: i + 100 async for i in f(range(5)) if 0 < i < 4}
+        self.assertEqual(
+            run_async(run_dict()),
+            ([], {11: 101, 12: 102, 13: 103}))
+
+        async def run_gen():
+            gen = (i + 10 async for i in f(range(5)) if 0 < i < 4)
+            return [g + 100 async for g in gen]
+        self.assertEqual(
+            run_async(run_gen()),
+            ([], [111, 112, 113]))
+
     def test_comp_5(self):
         async def f(it):
             for i in it:

Lib/test/test_sys_settrace.py

@@ -33,6 +33,10 @@ class asynctracecontext:
     async def __aexit__(self, *exc_info):
         self.output.append(-self.value)
 
+async def asynciter(iterable):
+    """Convert an iterable to an asynchronous iterator."""
+    for x in iterable:
+        yield x
 
 
 # A very basic example.  If this fails, we're in deep trouble.
@@ -721,6 +725,23 @@ class JumpTestCase(unittest.TestCase):
             output.append(6)
         output.append(7)
 
+    @async_jump_test(4, 5, [3, 5])
+    async def test_jump_out_of_async_for_block_forwards(output):
+        for i in [1]:
+            async for i in asynciter([1, 2]):
+                output.append(3)
+                output.append(4)
+            output.append(5)
+
+    @async_jump_test(5, 2, [2, 4, 2, 4, 5, 6])
+    async def test_jump_out_of_async_for_block_backwards(output):
+        for i in [1]:
+            output.append(2)
+            async for i in asynciter([1]):
+                output.append(4)
+                output.append(5)
+            output.append(6)
+
     @jump_test(1, 2, [3])
     def test_jump_to_codeless_line(output):
         output.append(1)
@@ -1000,6 +1021,17 @@ class JumpTestCase(unittest.TestCase):
             output.append(7)
         output.append(8)
 
+    @async_jump_test(1, 7, [7, 8])
+    async def test_jump_over_async_for_block_before_else(output):
+        output.append(1)
+        if not output:  # always false
+            async for i in asynciter([3]):
+                output.append(4)
+        else:
+            output.append(6)
+            output.append(7)
+        output.append(8)
+
     # The second set of 'jump' tests are for things that are not allowed:
 
     @jump_test(2, 3, [1], (ValueError, 'after'))
@@ -1051,12 +1083,24 @@ class JumpTestCase(unittest.TestCase):
         for i in 1, 2:
             output.append(3)
 
+    @async_jump_test(1, 3, [], (ValueError, 'into'))
+    async def test_no_jump_forwards_into_async_for_block(output):
+        output.append(1)
+        async for i in asynciter([1, 2]):
+            output.append(3)
+
     @jump_test(3, 2, [2, 2], (ValueError, 'into'))
     def test_no_jump_backwards_into_for_block(output):
         for i in 1, 2:
             output.append(2)
         output.append(3)
 
+    @async_jump_test(3, 2, [2, 2], (ValueError, 'into'))
+    async def test_no_jump_backwards_into_async_for_block(output):
+        async for i in asynciter([1, 2]):
+            output.append(2)
+        output.append(3)
+
     @jump_test(2, 4, [], (ValueError, 'into'))
     def test_no_jump_forwards_into_while_block(output):
         i = 1
@@ -1196,6 +1240,17 @@ class JumpTestCase(unittest.TestCase):
             output.append(7)
         output.append(8)
 
+    @async_jump_test(7, 4, [1, 6], (ValueError, 'into'))
+    async def test_no_jump_into_async_for_block_before_else(output):
+        output.append(1)
+        if not output:  # always false
+            async for i in asynciter([3]):
+                output.append(4)
+        else:
+            output.append(6)
+            output.append(7)
+        output.append(8)
+
     def test_no_jump_to_non_integers(self):
         self.run_test(no_jump_to_non_integers, 2, "Spam", [True])
 

Misc/NEWS.d/next/Core and Builtins/2018-03-18-13-56-14.bpo-33041.XwPhI2.rst

@@ -0,0 +1 @@
+Fixed jumping when the function contains an ``async for`` loop.

Python/compile.c

@@ -2383,24 +2383,19 @@ compiler_async_for(struct compiler *c, stmt_ty s)
     ADDOP(c, DUP_TOP);
     ADDOP_O(c, LOAD_GLOBAL, stop_aiter_error, names);
     ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
-    ADDOP_JABS(c, POP_JUMP_IF_FALSE, try_cleanup);
+    ADDOP_JABS(c, POP_JUMP_IF_TRUE, try_cleanup);
-
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
-    ADDOP(c, POP_TOP); /* for correct calculation of stack effect */
-    ADDOP(c, POP_BLOCK); /* for SETUP_LOOP */
-    ADDOP_JABS(c, JUMP_ABSOLUTE, after_loop_else);
-
-
-    compiler_use_next_block(c, try_cleanup);
     ADDOP(c, END_FINALLY);
 
     compiler_use_next_block(c, after_try);
     VISIT_SEQ(c, stmt, s->v.AsyncFor.body);
     ADDOP_JABS(c, JUMP_ABSOLUTE, try);
 
+    compiler_use_next_block(c, try_cleanup);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
+    ADDOP(c, POP_TOP); /* for correct calculation of stack effect */
     ADDOP(c, POP_BLOCK); /* for SETUP_LOOP */
     compiler_pop_fblock(c, LOOP, try);
 
@@ -3890,7 +3885,7 @@ compiler_async_comprehension_generator(struct compiler *c,
     _Py_IDENTIFIER(StopAsyncIteration);
 
     comprehension_ty gen;
-    basicblock *anchor, *if_cleanup, *try,
+    basicblock *if_cleanup, *try,
                *after_try, *except, *try_cleanup;
     Py_ssize_t i, n;
 
@@ -3901,12 +3896,11 @@ compiler_async_comprehension_generator(struct compiler *c,
 
     try = compiler_new_block(c);
     after_try = compiler_new_block(c);
-    try_cleanup = compiler_new_block(c);
     except = compiler_new_block(c);
     if_cleanup = compiler_new_block(c);
-    anchor = compiler_new_block(c);
+    try_cleanup = compiler_new_block(c);
 
-    if (if_cleanup == NULL || anchor == NULL ||
+    if (if_cleanup == NULL ||
             try == NULL || after_try == NULL ||
             except == NULL || try_cleanup == NULL) {
         return 0;
@@ -3945,16 +3939,7 @@ compiler_async_comprehension_generator(struct compiler *c,
     ADDOP(c, DUP_TOP);
     ADDOP_O(c, LOAD_GLOBAL, stop_aiter_error, names);
     ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
-    ADDOP_JABS(c, POP_JUMP_IF_FALSE, try_cleanup);
+    ADDOP_JABS(c, POP_JUMP_IF_TRUE, try_cleanup);
-
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
-    ADDOP_JABS(c, JUMP_ABSOLUTE, anchor);
-
-
-    compiler_use_next_block(c, try_cleanup);
     ADDOP(c, END_FINALLY);
 
     compiler_use_next_block(c, after_try);
@@ -4003,7 +3988,12 @@ compiler_async_comprehension_generator(struct compiler *c,
     }
     compiler_use_next_block(c, if_cleanup);
     ADDOP_JABS(c, JUMP_ABSOLUTE, try);
-    compiler_use_next_block(c, anchor);
+
+    compiler_use_next_block(c, try_cleanup);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
     ADDOP(c, POP_TOP);
 
     return 1;