Skip to content

Commit 175cfd0

Browse files
committed
Fix list comprehensions raising StopIteration when empty due to condition
1 parent ef751a5 commit 175cfd0

File tree

2 files changed

+33
-12
lines changed

2 files changed

+33
-12
lines changed

graalpython/com.oracle.graal.python.test/src/com/oracle/graal/python/test/generator/BuiltinIntrinsificationTests.java

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2017, 2018, Oracle and/or its affiliates.
2+
* Copyright (c) 2017, 2020, Oracle and/or its affiliates.
33
* Copyright (c) 2014, Regents of the University of California
44
*
55
* All rights reserved.
@@ -38,4 +38,18 @@ public void simpleListComp() {
3838
"print(ll)";
3939
assertPrints("[0, 1, 2, 3, 4]\n", source);
4040
}
41+
42+
@Test
43+
public void simpleListCompCondition() {
44+
String source = "ll = list(i for i in range(5) if i % 2 == 0)\n" + //
45+
"print(ll)";
46+
assertPrints("[0, 2, 4]\n", source);
47+
}
48+
49+
@Test
50+
public void emptyListComp() {
51+
String source = "ll = list(i for i in range(5) if i == 6)\n" + //
52+
"print(ll)";
53+
assertPrints("[]\n", source);
54+
}
4155
}

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/list/ListBuiltins.java

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -230,17 +230,24 @@ PNone listPGenerator(VirtualFrame frame, PList list, PGenerator iterable,
230230
SequenceStorage storage = EmptySequenceStorage.INSTANCE;
231231

232232
PRangeIterator range = (PRangeIterator) iterable.getIterator();
233-
final int len = range.getLength(stepProfile, positiveRangeProfile);
234-
if (len > 0) {
235-
Object value = getNextNode.execute(frame, iterObj);
236-
storage = SequenceStorageFactory.createStorage(value, len);
237-
storage = appendNode.execute(storage, value, ListGeneralizationNode.SUPPLIER);
238-
while (true) {
239-
try {
240-
storage = appendNode.execute(storage, getNextNode.execute(frame, iterObj), ListGeneralizationNode.SUPPLIER);
241-
} catch (PException e) {
242-
e.expectStopIteration(errorProfile);
243-
break;
233+
final int estimatedMaxLen = range.getLength(stepProfile, positiveRangeProfile);
234+
if (estimatedMaxLen > 0) {
235+
Object value = null;
236+
try {
237+
value = getNextNode.execute(frame, iterObj);
238+
} catch (PException e) {
239+
e.expectStopIteration(errorProfile);
240+
}
241+
if (value != null) {
242+
storage = SequenceStorageFactory.createStorage(value, estimatedMaxLen);
243+
storage = appendNode.execute(storage, value, ListGeneralizationNode.SUPPLIER);
244+
while (true) {
245+
try {
246+
storage = appendNode.execute(storage, getNextNode.execute(frame, iterObj), ListGeneralizationNode.SUPPLIER);
247+
} catch (PException e) {
248+
e.expectStopIteration(errorProfile);
249+
break;
250+
}
244251
}
245252
}
246253
}

0 commit comments

Comments
 (0)