Skip to content

Commit e37dfb1

Browse files
authored
[mypyc] Fix evaluation of iterable in list comprehension twice (#10599)
This could result in a crash if the second evaluation results in a shorter list, such as in this example (besides being incorrect overall): ``` a = [s for s in f.readlines()] ``` `f.readlines()` was called twice, resulting in an empty list on the second call. This caused the list object constructed in the comprehension to have a NULL item, which is invalid. This fixes `mypy --install-types` in compiled mode. Fixes #10596.
1 parent 0a6a48c commit e37dfb1

File tree

2 files changed

+23
-7
lines changed

2 files changed

+23
-7
lines changed

mypyc/irbuild/for_helpers.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,10 @@ def for_loop_helper(builder: IRBuilder, index: Lvalue, expr: Expression,
8888
builder.activate_block(exit_block)
8989

9090

91-
def for_loop_helper_with_index(builder: IRBuilder, index: Lvalue, expr: Expression,
91+
def for_loop_helper_with_index(builder: IRBuilder,
92+
index: Lvalue,
93+
expr: Expression,
94+
expr_reg: Value,
9295
body_insts: Callable[[Value], None], line: int) -> None:
9396
"""Generate IR for a sequence iteration.
9497
@@ -101,7 +104,6 @@ def for_loop_helper_with_index(builder: IRBuilder, index: Lvalue, expr: Expressi
101104
body_insts: a function that generates the body of the loop.
102105
It needs a index as parameter.
103106
"""
104-
expr_reg = builder.accept(expr)
105107
assert is_sequence_rprimitive(expr_reg.type)
106108
target_type = builder.get_sequence_type(expr)
107109

@@ -163,16 +165,15 @@ def sequence_from_generator_preallocate_helper(
163165
if len(gen.sequences) == 1 and len(gen.indices) == 1 and len(gen.condlists[0]) == 0:
164166
rtype = builder.node_type(gen.sequences[0])
165167
if is_list_rprimitive(rtype) or is_tuple_rprimitive(rtype):
166-
167-
length = builder.builder.builtin_len(builder.accept(gen.sequences[0]),
168-
gen.line, use_pyssize_t=True)
168+
sequence = builder.accept(gen.sequences[0])
169+
length = builder.builder.builtin_len(sequence, gen.line, use_pyssize_t=True)
169170
target_op = empty_op_llbuilder(length, gen.line)
170171

171172
def set_item(item_index: Value) -> None:
172173
e = builder.accept(gen.left_expr)
173174
builder.call_c(set_item_op, [target_op, item_index, e], gen.line)
174175

175-
for_loop_helper_with_index(builder, gen.indices[0], gen.sequences[0],
176+
for_loop_helper_with_index(builder, gen.indices[0], gen.sequences[0], sequence,
176177
set_item, gen.line)
177178

178179
return target_op

mypyc/test-data/run-misc.test

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -485,6 +485,8 @@ print(native.x)
485485
77
486486

487487
[case testComprehensions]
488+
from typing import List
489+
488490
# A list comprehension
489491
l = [str(x) + " " + str(y) + " " + str(x*y) for x in range(10)
490492
if x != 6 if x != 5 for y in range(x) if y*x != 8]
@@ -498,6 +500,17 @@ def pred(x: int) -> bool:
498500
# eventually and will raise an exception.
499501
l2 = [x for x in range(10) if x <= 6 if pred(x)]
500502

503+
src = ['x']
504+
505+
def f() -> List[str]:
506+
global src
507+
res = src
508+
src = []
509+
return res
510+
511+
l3 = [s for s in f()]
512+
l4 = [s for s in f()]
513+
501514
# A dictionary comprehension
502515
d = {k: k*k for k in range(10) if k != 5 if k != 6}
503516

@@ -506,10 +519,12 @@ s = {str(x) + " " + str(y) + " " + str(x*y) for x in range(10)
506519
if x != 6 if x != 5 for y in range(x) if y*x != 8}
507520

508521
[file driver.py]
509-
from native import l, l2, d, s
522+
from native import l, l2, l3, l4, d, s
510523
for a in l:
511524
print(a)
512525
print(tuple(l2))
526+
assert l3 == ['x']
527+
assert l4 == []
513528
for k in sorted(d):
514529
print(k, d[k])
515530
for a in sorted(s):

0 commit comments

Comments
 (0)