Allow if filtering in list comprehensions -- MOS_MIGRATED_REVID=93881507
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/LValue.java b/src/main/java/com/google/devtools/build/lib/syntax/LValue.java index 846c5b5..fcaf3a8 100644 --- a/src/main/java/com/google/devtools/build/lib/syntax/LValue.java +++ b/src/main/java/com/google/devtools/build/lib/syntax/LValue.java
@@ -29,7 +29,6 @@ * An LValue can be a simple variable or something more complex like a tuple. */ public class LValue implements Serializable { - // Currently, expr can only be an Ident, but we plan to support more. private final Expression expr; public LValue(Expression expr) {
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/ListComprehension.java b/src/main/java/com/google/devtools/build/lib/syntax/ListComprehension.java index 5d3b655..bd83e21 100644 --- a/src/main/java/com/google/devtools/build/lib/syntax/ListComprehension.java +++ b/src/main/java/com/google/devtools/build/lib/syntax/ListComprehension.java
@@ -26,12 +26,12 @@ /** * Syntax node for lists comprehension expressions. * - * A list comprehension contains one or more clauses, e.g. + * <p> A list comprehension contains one or more clauses, e.g. * [a+d for a in b if c for d in e] * contains three clauses: "for a in b", "if c", "for d in e". * For and If clauses can happen in any order, except that the first one has to be a For. * - * The code above can be expanded as: + * <p> The code above can be expanded as: * <pre> * for a in b: * if c: @@ -78,8 +78,6 @@ public abstract Expression getExpression(); } - // TODO(bazel-team): Support IfClause - /** * A for clause in a list comprehension, e.g. "for a in b" in the example above. */ @@ -126,6 +124,45 @@ } } + /** + * A if clause in a list comprehension, e.g. "if c" in the example above. + */ + public final class IfClause implements Clause { + private final Expression condition; + + public IfClause(Expression condition) { + this.condition = condition; + } + + @Override + public void eval(Environment env, List<Object> result, int step) + throws EvalException, InterruptedException { + if (EvalUtils.toBoolean(condition.eval(env))) { + evalStep(env, result, step); + } + } + + @Override + public void validate(ValidationEnvironment env) throws EvalException { + condition.validate(env); + } + + @Override + public LValue getLValue() { + return null; + } + + @Override + public Expression getExpression() { + return condition; + } + + @Override + public String toString() { + return String.format("if %s", condition); + } + } + private List<Clause> clauses; /** The return expression, e.g. "a+d" in the example above */ private final Expression elementExpression; @@ -156,11 +193,19 @@ * TODO(bazel-team): Remove this side-effect. Clauses should be passed to the constructor * instead. */ - public void add(Expression loopVar, Expression listExpression) { + void addFor(Expression loopVar, Expression listExpression) { Clause forClause = new ForClause(new LValue(loopVar), listExpression); clauses.add(forClause); } + /** + * Add a new ForClause to the list comprehension. + * TODO(bazel-team): Remove this side-effect. + */ + void addIf(Expression condition) { + clauses.add(new IfClause(condition)); + } + public List<Clause> getClauses() { return Collections.unmodifiableList(clauses); } @@ -191,7 +236,7 @@ * recursively call evalStep any number of times. After the last clause, * elementExpression is evaluated and added to the results. * - * In the expanded example above, you can consider that evalStep is equivalent to + * <p> In the expanded example above, you can consider that evalStep is equivalent to * evaluating the line number step. */ private void evalStep(Environment env, List<Object> result, int step)
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Parser.java b/src/main/java/com/google/devtools/build/lib/syntax/Parser.java index 6381c15..ae267e5 100644 --- a/src/main/java/com/google/devtools/build/lib/syntax/Parser.java +++ b/src/main/java/com/google/devtools/build/lib/syntax/Parser.java
@@ -786,15 +786,15 @@ nextToken(); Expression loopVar = parseForLoopVariables(); expect(TokenKind.IN); - Expression listExpression = parseExpression(); - listComprehension.add(loopVar, listExpression); + // The expression cannot be a ternary expression ('x if y else z') due to + // conflicts in Python grammar ('if' is used by the comprehension). + Expression listExpression = parseNonTupleExpression(0); + listComprehension.addFor(loopVar, listExpression); break; case IF: - reportError(lexer.createLocation(token.left, token.right), - "List comprehension with filtering is not yet supported"); nextToken(); - parseExpression(); // condition + listComprehension.addIf(parseExpression()); break; case RBRACKET:
diff --git a/src/test/java/com/google/devtools/build/lib/syntax/EvaluationTest.java b/src/test/java/com/google/devtools/build/lib/syntax/EvaluationTest.java index b410fa0..b3e5bcf 100644 --- a/src/test/java/com/google/devtools/build/lib/syntax/EvaluationTest.java +++ b/src/test/java/com/google/devtools/build/lib/syntax/EvaluationTest.java
@@ -394,6 +394,29 @@ } @Test + public void testListComprehensionsWithFiltering() throws Exception { + eval("range3 = [0, 1, 2]"); // used below + + assertThat(eval("[a for a in (4, None, 2, None, 1) if a != None]").toString()) + .isEqualTo("[4, 2, 1]"); + assertThat(eval("[b+c for b in [0, 1, 2] for c in [0, 1, 2] if b + c > 2]").toString()) + .isEqualTo("[3, 3, 4]"); + assertThat(eval("[d+e for d in range3 if d % 2 == 1 for e in range3]").toString()) + .isEqualTo("[1, 2, 3]"); + assertThat(eval( + "[[f,g] for f in [0, 1, 2, 3, 4] if f for g in [5, 6, 7, 8] if f * g % 12 == 0 ]") + .toString()).isEqualTo("[[2, 6], [3, 8], [4, 6]]"); + assertThat(eval("[h for h in [4, 2, 0, 1] if h]").toString()) + .isEqualTo("[4, 2, 1]"); + } + + @Test + public void testListComprehensionDefinitionOrder() throws Exception { + checkEvalErrorContains("name 'y' is not defined", + "[x for x in (1, 2) if y for y in (3, 4)]"); + } + + @Test public void testTupleDestructuring() throws Exception { eval("a, b = 1, 2"); assertThat(lookup("a")).isEqualTo(1);