Skip to content

Commit b9744e9

Browse files
bpo-33041: Fixed jumping if the function contains an "async for" loop. (GH-6154)
1 parent c71edab commit b9744e9

4 files changed

Lines changed: 103 additions & 27 deletions

File tree

Lib/test/test_coroutines.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1846,6 +1846,36 @@ async def run_gen():
18461846
run_async(run_gen()),
18471847
([], [121]))
18481848

1849+
def test_comp_4_2(self):
1850+
async def f(it):
1851+
for i in it:
1852+
yield i
1853+
1854+
async def run_list():
1855+
return [i + 10 async for i in f(range(5)) if 0 < i < 4]
1856+
self.assertEqual(
1857+
run_async(run_list()),
1858+
([], [11, 12, 13]))
1859+
1860+
async def run_set():
1861+
return {i + 10 async for i in f(range(5)) if 0 < i < 4}
1862+
self.assertEqual(
1863+
run_async(run_set()),
1864+
([], {11, 12, 13}))
1865+
1866+
async def run_dict():
1867+
return {i + 10: i + 100 async for i in f(range(5)) if 0 < i < 4}
1868+
self.assertEqual(
1869+
run_async(run_dict()),
1870+
([], {11: 101, 12: 102, 13: 103}))
1871+
1872+
async def run_gen():
1873+
gen = (i + 10 async for i in f(range(5)) if 0 < i < 4)
1874+
return [g + 100 async for g in gen]
1875+
self.assertEqual(
1876+
run_async(run_gen()),
1877+
([], [111, 112, 113]))
1878+
18491879
def test_comp_5(self):
18501880
async def f(it):
18511881
for i in it:

Lib/test/test_sys_settrace.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,10 @@ async def __aenter__(self):
3333
async def __aexit__(self, *exc_info):
3434
self.output.append(-self.value)
3535

36+
async def asynciter(iterable):
37+
"""Convert an iterable to an asynchronous iterator."""
38+
for x in iterable:
39+
yield x
3640

3741

3842
# A very basic example. If this fails, we're in deep trouble.
@@ -721,6 +725,23 @@ def test_jump_out_of_block_backwards(output):
721725
output.append(6)
722726
output.append(7)
723727

728+
@async_jump_test(4, 5, [3, 5])
729+
async def test_jump_out_of_async_for_block_forwards(output):
730+
for i in [1]:
731+
async for i in asynciter([1, 2]):
732+
output.append(3)
733+
output.append(4)
734+
output.append(5)
735+
736+
@async_jump_test(5, 2, [2, 4, 2, 4, 5, 6])
737+
async def test_jump_out_of_async_for_block_backwards(output):
738+
for i in [1]:
739+
output.append(2)
740+
async for i in asynciter([1]):
741+
output.append(4)
742+
output.append(5)
743+
output.append(6)
744+
724745
@jump_test(1, 2, [3])
725746
def test_jump_to_codeless_line(output):
726747
output.append(1)
@@ -1000,6 +1021,17 @@ def test_jump_over_for_block_before_else(output):
10001021
output.append(7)
10011022
output.append(8)
10021023

1024+
@async_jump_test(1, 7, [7, 8])
1025+
async def test_jump_over_async_for_block_before_else(output):
1026+
output.append(1)
1027+
if not output: # always false
1028+
async for i in asynciter([3]):
1029+
output.append(4)
1030+
else:
1031+
output.append(6)
1032+
output.append(7)
1033+
output.append(8)
1034+
10031035
# The second set of 'jump' tests are for things that are not allowed:
10041036

10051037
@jump_test(2, 3, [1], (ValueError, 'after'))
@@ -1051,12 +1083,24 @@ def test_no_jump_forwards_into_for_block(output):
10511083
for i in 1, 2:
10521084
output.append(3)
10531085

1086+
@async_jump_test(1, 3, [], (ValueError, 'into'))
1087+
async def test_no_jump_forwards_into_async_for_block(output):
1088+
output.append(1)
1089+
async for i in asynciter([1, 2]):
1090+
output.append(3)
1091+
10541092
@jump_test(3, 2, [2, 2], (ValueError, 'into'))
10551093
def test_no_jump_backwards_into_for_block(output):
10561094
for i in 1, 2:
10571095
output.append(2)
10581096
output.append(3)
10591097

1098+
@async_jump_test(3, 2, [2, 2], (ValueError, 'into'))
1099+
async def test_no_jump_backwards_into_async_for_block(output):
1100+
async for i in asynciter([1, 2]):
1101+
output.append(2)
1102+
output.append(3)
1103+
10601104
@jump_test(2, 4, [], (ValueError, 'into'))
10611105
def test_no_jump_forwards_into_while_block(output):
10621106
i = 1
@@ -1196,6 +1240,17 @@ def test_no_jump_into_for_block_before_else(output):
11961240
output.append(7)
11971241
output.append(8)
11981242

1243+
@async_jump_test(7, 4, [1, 6], (ValueError, 'into'))
1244+
async def test_no_jump_into_async_for_block_before_else(output):
1245+
output.append(1)
1246+
if not output: # always false
1247+
async for i in asynciter([3]):
1248+
output.append(4)
1249+
else:
1250+
output.append(6)
1251+
output.append(7)
1252+
output.append(8)
1253+
11991254
def test_no_jump_to_non_integers(self):
12001255
self.run_test(no_jump_to_non_integers, 2, "Spam", [True])
12011256

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Fixed jumping when the function contains an ``async for`` loop.

Python/compile.c

Lines changed: 17 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2383,24 +2383,19 @@ compiler_async_for(struct compiler *c, stmt_ty s)
23832383
ADDOP(c, DUP_TOP);
23842384
ADDOP_O(c, LOAD_GLOBAL, stop_aiter_error, names);
23852385
ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
2386-
ADDOP_JABS(c, POP_JUMP_IF_FALSE, try_cleanup);
2387-
2388-
ADDOP(c, POP_TOP);
2389-
ADDOP(c, POP_TOP);
2390-
ADDOP(c, POP_TOP);
2391-
ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
2392-
ADDOP(c, POP_TOP); /* for correct calculation of stack effect */
2393-
ADDOP(c, POP_BLOCK); /* for SETUP_LOOP */
2394-
ADDOP_JABS(c, JUMP_ABSOLUTE, after_loop_else);
2395-
2396-
2397-
compiler_use_next_block(c, try_cleanup);
2386+
ADDOP_JABS(c, POP_JUMP_IF_TRUE, try_cleanup);
23982387
ADDOP(c, END_FINALLY);
23992388

24002389
compiler_use_next_block(c, after_try);
24012390
VISIT_SEQ(c, stmt, s->v.AsyncFor.body);
24022391
ADDOP_JABS(c, JUMP_ABSOLUTE, try);
24032392

2393+
compiler_use_next_block(c, try_cleanup);
2394+
ADDOP(c, POP_TOP);
2395+
ADDOP(c, POP_TOP);
2396+
ADDOP(c, POP_TOP);
2397+
ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
2398+
ADDOP(c, POP_TOP); /* for correct calculation of stack effect */
24042399
ADDOP(c, POP_BLOCK); /* for SETUP_LOOP */
24052400
compiler_pop_fblock(c, LOOP, try);
24062401

@@ -3890,7 +3885,7 @@ compiler_async_comprehension_generator(struct compiler *c,
38903885
_Py_IDENTIFIER(StopAsyncIteration);
38913886

38923887
comprehension_ty gen;
3893-
basicblock *anchor, *if_cleanup, *try,
3888+
basicblock *if_cleanup, *try,
38943889
*after_try, *except, *try_cleanup;
38953890
Py_ssize_t i, n;
38963891

@@ -3901,12 +3896,11 @@ compiler_async_comprehension_generator(struct compiler *c,
39013896

39023897
try = compiler_new_block(c);
39033898
after_try = compiler_new_block(c);
3904-
try_cleanup = compiler_new_block(c);
39053899
except = compiler_new_block(c);
39063900
if_cleanup = compiler_new_block(c);
3907-
anchor = compiler_new_block(c);
3901+
try_cleanup = compiler_new_block(c);
39083902

3909-
if (if_cleanup == NULL || anchor == NULL ||
3903+
if (if_cleanup == NULL ||
39103904
try == NULL || after_try == NULL ||
39113905
except == NULL || try_cleanup == NULL) {
39123906
return 0;
@@ -3945,16 +3939,7 @@ compiler_async_comprehension_generator(struct compiler *c,
39453939
ADDOP(c, DUP_TOP);
39463940
ADDOP_O(c, LOAD_GLOBAL, stop_aiter_error, names);
39473941
ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
3948-
ADDOP_JABS(c, POP_JUMP_IF_FALSE, try_cleanup);
3949-
3950-
ADDOP(c, POP_TOP);
3951-
ADDOP(c, POP_TOP);
3952-
ADDOP(c, POP_TOP);
3953-
ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
3954-
ADDOP_JABS(c, JUMP_ABSOLUTE, anchor);
3955-
3956-
3957-
compiler_use_next_block(c, try_cleanup);
3942+
ADDOP_JABS(c, POP_JUMP_IF_TRUE, try_cleanup);
39583943
ADDOP(c, END_FINALLY);
39593944

39603945
compiler_use_next_block(c, after_try);
@@ -4003,7 +3988,12 @@ compiler_async_comprehension_generator(struct compiler *c,
40033988
}
40043989
compiler_use_next_block(c, if_cleanup);
40053990
ADDOP_JABS(c, JUMP_ABSOLUTE, try);
4006-
compiler_use_next_block(c, anchor);
3991+
3992+
compiler_use_next_block(c, try_cleanup);
3993+
ADDOP(c, POP_TOP);
3994+
ADDOP(c, POP_TOP);
3995+
ADDOP(c, POP_TOP);
3996+
ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
40073997
ADDOP(c, POP_TOP);
40083998

40093999
return 1;

0 commit comments

Comments
 (0)