diff --git a/python-squid/src/main/java/org/sonar/python/api/tree/PyListOrSetCompExpressionTree.java b/python-squid/src/main/java/org/sonar/python/api/tree/PyComprehensionExpressionTree.java similarity index 93% rename from python-squid/src/main/java/org/sonar/python/api/tree/PyListOrSetCompExpressionTree.java rename to python-squid/src/main/java/org/sonar/python/api/tree/PyComprehensionExpressionTree.java index 51512e9a2a..065963809e 100644 --- a/python-squid/src/main/java/org/sonar/python/api/tree/PyListOrSetCompExpressionTree.java +++ b/python-squid/src/main/java/org/sonar/python/api/tree/PyComprehensionExpressionTree.java @@ -19,7 +19,7 @@ */ package org.sonar.python.api.tree; -public interface PyListOrSetCompExpressionTree extends PyExpressionTree { +public interface PyComprehensionExpressionTree extends PyExpressionTree { PyExpressionTree resultExpression(); diff --git a/python-squid/src/main/java/org/sonar/python/api/tree/PyTreeVisitor.java b/python-squid/src/main/java/org/sonar/python/api/tree/PyTreeVisitor.java index bac5fe11b3..fe06c66e95 100644 --- a/python-squid/src/main/java/org/sonar/python/api/tree/PyTreeVisitor.java +++ b/python-squid/src/main/java/org/sonar/python/api/tree/PyTreeVisitor.java @@ -137,7 +137,7 @@ public interface PyTreeVisitor { void visitConditionalExpression(PyConditionalExpressionTree pyConditionalExpressionTree); - void visitPyListOrSetCompExpression(PyListOrSetCompExpressionTree tree); + void visitPyListOrSetCompExpression(PyComprehensionExpressionTree tree); void visitComprehensionFor(PyComprehensionForTree tree); diff --git a/python-squid/src/main/java/org/sonar/python/api/tree/Tree.java b/python-squid/src/main/java/org/sonar/python/api/tree/Tree.java index 43cacad9f0..c7e1b6bb0b 100644 --- a/python-squid/src/main/java/org/sonar/python/api/tree/Tree.java +++ b/python-squid/src/main/java/org/sonar/python/api/tree/Tree.java @@ -162,8 +162,9 @@ enum Kind { TUPLE(PyTupleTree.class), DICT_COMPREHENSION(PyDictCompExpressionTree.class), - LIST_COMPREHENSION(PyListOrSetCompExpressionTree.class), - SET_COMPREHENSION(PyListOrSetCompExpressionTree.class), + LIST_COMPREHENSION(PyComprehensionExpressionTree.class), + SET_COMPREHENSION(PyComprehensionExpressionTree.class), + GENERATOR_EXPR(PyComprehensionExpressionTree.class), COMP_FOR(PyComprehensionForTree.class), COMP_IF(PyComprehensionIfTree.class), diff --git a/python-squid/src/main/java/org/sonar/python/tree/BaseTreeVisitor.java b/python-squid/src/main/java/org/sonar/python/tree/BaseTreeVisitor.java index 28a9faeede..a6beed02f9 100644 --- a/python-squid/src/main/java/org/sonar/python/tree/BaseTreeVisitor.java +++ b/python-squid/src/main/java/org/sonar/python/tree/BaseTreeVisitor.java @@ -57,7 +57,7 @@ import org.sonar.python.api.tree.PyKeyValuePairTree; import org.sonar.python.api.tree.PyLambdaExpressionTree; import org.sonar.python.api.tree.PyListLiteralTree; -import org.sonar.python.api.tree.PyListOrSetCompExpressionTree; +import org.sonar.python.api.tree.PyComprehensionExpressionTree; import org.sonar.python.api.tree.PyNameTree; import org.sonar.python.api.tree.PyNoneExpressionTree; import org.sonar.python.api.tree.PyNonlocalStatementTree; @@ -437,7 +437,7 @@ public void visitConditionalExpression(PyConditionalExpressionTree pyConditional } @Override - public void visitPyListOrSetCompExpression(PyListOrSetCompExpressionTree tree) { + public void visitPyListOrSetCompExpression(PyComprehensionExpressionTree tree) { scan(tree.resultExpression()); scan(tree.comprehensionFor()); } diff --git a/python-squid/src/main/java/org/sonar/python/tree/PyListOrSetCompExpressionTreeImpl.java b/python-squid/src/main/java/org/sonar/python/tree/PyComprehensionExpressionTreeImpl.java similarity index 89% rename from python-squid/src/main/java/org/sonar/python/tree/PyListOrSetCompExpressionTreeImpl.java rename to python-squid/src/main/java/org/sonar/python/tree/PyComprehensionExpressionTreeImpl.java index b4558be977..22a7c0ffd4 100644 --- a/python-squid/src/main/java/org/sonar/python/tree/PyListOrSetCompExpressionTreeImpl.java +++ b/python-squid/src/main/java/org/sonar/python/tree/PyComprehensionExpressionTreeImpl.java @@ -24,17 +24,17 @@ import java.util.List; import org.sonar.python.api.tree.PyComprehensionForTree; import org.sonar.python.api.tree.PyExpressionTree; -import org.sonar.python.api.tree.PyListOrSetCompExpressionTree; +import org.sonar.python.api.tree.PyComprehensionExpressionTree; import org.sonar.python.api.tree.PyTreeVisitor; import org.sonar.python.api.tree.Tree; -public class PyListOrSetCompExpressionTreeImpl extends PyTree implements PyListOrSetCompExpressionTree { +public class PyComprehensionExpressionTreeImpl extends PyTree implements PyComprehensionExpressionTree { private final Kind kind; private final PyExpressionTree resultExpression; private final PyComprehensionForTree comprehensionFor; - public PyListOrSetCompExpressionTreeImpl(Kind kind, Token openingToken, PyExpressionTree resultExpression, + public PyComprehensionExpressionTreeImpl(Kind kind, Token openingToken, PyExpressionTree resultExpression, PyComprehensionForTree compFor, Token closingToken) { super(openingToken, closingToken); this.kind = kind; diff --git a/python-squid/src/main/java/org/sonar/python/tree/PythonTreeMaker.java b/python-squid/src/main/java/org/sonar/python/tree/PythonTreeMaker.java index 7aa629c1a0..18032d9269 100644 --- a/python-squid/src/main/java/org/sonar/python/tree/PythonTreeMaker.java +++ b/python-squid/src/main/java/org/sonar/python/tree/PythonTreeMaker.java @@ -762,7 +762,7 @@ private PyExpressionTree dictOrSetLiteral(AstNode astNode) { return new PyDictCompExpressionTreeImpl(lCurlyBrace, keyExpression, colon.getToken(), valueExpression, compFor, rCurlyBrace); } else { PyExpressionTree resultExpression = expression(dictOrSetMaker.getFirstChild(PythonGrammar.TEST, PythonGrammar.STAR_EXPR)); - return new PyListOrSetCompExpressionTreeImpl(Tree.Kind.SET_COMPREHENSION, lCurlyBrace, resultExpression, compFor, rCurlyBrace); + return new PyComprehensionExpressionTreeImpl(Tree.Kind.SET_COMPREHENSION, lCurlyBrace, resultExpression, compFor, rCurlyBrace); } } List commas = dictOrSetMaker.getChildren(PythonPunctuator.COMMA).stream().map(AstNode::getToken).collect(Collectors.toList()); @@ -800,7 +800,10 @@ private PyExpressionTree parenthesized(AstNode atom) { return new PyTupleTreeImpl(atom, lPar, Collections.emptyList(), Collections.emptyList(), rPar); } - // TODO COMP_FOR + AstNode compFor = testListComp.getFirstChild(PythonGrammar.COMP_FOR); + if (compFor != null) { + return new PyComprehensionExpressionTreeImpl(Tree.Kind.GENERATOR_EXPR, lPar, expression(testListComp.getFirstChild()), compFor(compFor), rPar); + } PyExpressionListTree expressionList = expressionList(testListComp); List commas = testListComp.getChildren(PythonPunctuator.COMMA); if (commas.isEmpty()) { @@ -914,7 +917,7 @@ private PyExpressionTree listLiteral(AstNode astNode) { AstNode compForNode = testListComp.getFirstChild(PythonGrammar.COMP_FOR); if (compForNode != null) { PyExpressionTree resultExpression = expression(testListComp.getFirstChild(PythonGrammar.TEST, PythonGrammar.STAR_EXPR)); - return new PyListOrSetCompExpressionTreeImpl(Tree.Kind.LIST_COMPREHENSION, leftBracket, resultExpression, compFor(compForNode), rightBracket); + return new PyComprehensionExpressionTreeImpl(Tree.Kind.LIST_COMPREHENSION, leftBracket, resultExpression, compFor(compForNode), rightBracket); } elements = expressionList(testListComp); } else { diff --git a/python-squid/src/test/java/org/sonar/python/tree/BaseTreeVisitorTest.java b/python-squid/src/test/java/org/sonar/python/tree/BaseTreeVisitorTest.java index da861699e9..c71ec8ab74 100644 --- a/python-squid/src/test/java/org/sonar/python/tree/BaseTreeVisitorTest.java +++ b/python-squid/src/test/java/org/sonar/python/tree/BaseTreeVisitorTest.java @@ -43,7 +43,7 @@ import org.sonar.python.api.tree.PyImportNameTree; import org.sonar.python.api.tree.PyLambdaExpressionTree; import org.sonar.python.api.tree.PyListLiteralTree; -import org.sonar.python.api.tree.PyListOrSetCompExpressionTree; +import org.sonar.python.api.tree.PyComprehensionExpressionTree; import org.sonar.python.api.tree.PyNameTree; import org.sonar.python.api.tree.PyNumericLiteralTree; import org.sonar.python.api.tree.PyParameterTree; @@ -305,7 +305,7 @@ public void cond_expression() { @Test public void list_or_set_comprehension() { setRootRule(PythonGrammar.EXPR); - PyListOrSetCompExpressionTree expr = (PyListOrSetCompExpressionTree) parse("[x+1 for x in [42, 43] if cond(x)]", treeMaker::expression); + PyComprehensionExpressionTree expr = (PyComprehensionExpressionTree) parse("[x+1 for x in [42, 43] if cond(x)]", treeMaker::expression); BaseTreeVisitor visitor = spy(BaseTreeVisitor.class); expr.accept(visitor); diff --git a/python-squid/src/test/java/org/sonar/python/tree/PythonTreeMakerTest.java b/python-squid/src/test/java/org/sonar/python/tree/PythonTreeMakerTest.java index cf80482316..f17f85a2b6 100644 --- a/python-squid/src/test/java/org/sonar/python/tree/PythonTreeMakerTest.java +++ b/python-squid/src/test/java/org/sonar/python/tree/PythonTreeMakerTest.java @@ -65,7 +65,7 @@ import org.sonar.python.api.tree.PyKeyValuePairTree; import org.sonar.python.api.tree.PyLambdaExpressionTree; import org.sonar.python.api.tree.PyListLiteralTree; -import org.sonar.python.api.tree.PyListOrSetCompExpressionTree; +import org.sonar.python.api.tree.PyComprehensionExpressionTree; import org.sonar.python.api.tree.PyNameTree; import org.sonar.python.api.tree.PyNoneExpressionTree; import org.sonar.python.api.tree.PyNonlocalStatementTree; @@ -1497,8 +1497,8 @@ public void list_literal() { @Test public void list_comprehension() { setRootRule(PythonGrammar.TEST); - PyListOrSetCompExpressionTree comprehension = - (PyListOrSetCompExpressionTree) parse("[x+y for x,y in [(42, 43)]]", treeMaker::expression); + PyComprehensionExpressionTree comprehension = + (PyComprehensionExpressionTree) parse("[x+y for x,y in [(42, 43)]]", treeMaker::expression); assertThat(comprehension.getKind()).isEqualTo(Tree.Kind.LIST_COMPREHENSION); assertThat(comprehension.resultExpression().getKind()).isEqualTo(Tree.Kind.PLUS); assertThat(comprehension.children()).hasSize(2); @@ -1515,8 +1515,8 @@ public void list_comprehension() { @Test public void list_comprehension_with_if() { setRootRule(PythonGrammar.TEST); - PyListOrSetCompExpressionTree comprehension = - (PyListOrSetCompExpressionTree) parse("[x+1 for x in [42, 43] if x%2==0]", treeMaker::expression); + PyComprehensionExpressionTree comprehension = + (PyComprehensionExpressionTree) parse("[x+1 for x in [42, 43] if x%2==0]", treeMaker::expression); assertThat(comprehension.getKind()).isEqualTo(Tree.Kind.LIST_COMPREHENSION); PyComprehensionForTree forClause = comprehension.comprehensionFor(); assertThat(forClause.nestedClause().getKind()).isEqualTo(Tree.Kind.COMP_IF); @@ -1530,8 +1530,8 @@ public void list_comprehension_with_if() { @Test public void list_comprehension_with_nested_for() { setRootRule(PythonGrammar.TEST); - PyListOrSetCompExpressionTree comprehension = - (PyListOrSetCompExpressionTree) parse("[x+y for x in [42, 43] for y in ('a', 0)]", treeMaker::expression); + PyComprehensionExpressionTree comprehension = + (PyComprehensionExpressionTree) parse("[x+y for x in [42, 43] for y in ('a', 0)]", treeMaker::expression); assertThat(comprehension.getKind()).isEqualTo(Tree.Kind.LIST_COMPREHENSION); PyComprehensionForTree forClause = comprehension.comprehensionFor(); assertThat(forClause.iterable().getKind()).isEqualTo(Tree.Kind.LIST_LITERAL); @@ -1552,6 +1552,19 @@ public void parenthesized_expression() { assertThat(parenthesized.expression().getKind()).isEqualTo(Tree.Kind.YIELD_EXPR); } + + @Test + public void generator_expression() { + setRootRule(PythonGrammar.TEST); + PyComprehensionExpressionTree generator = (PyComprehensionExpressionTree) parse("(x*x for x in range(10))", treeMaker::expression); + assertThat(generator.getKind()).isEqualTo(Tree.Kind.GENERATOR_EXPR); + assertThat(generator.children()).hasSize(2); + assertThat(generator.firstToken().getValue()).isEqualTo("("); + assertThat(generator.lastToken().getValue()).isEqualTo(")"); + assertThat(generator.resultExpression().getKind()).isEqualTo(Tree.Kind.MULTIPLICATION); + assertThat(generator.comprehensionFor().iterable().getKind()).isEqualTo(Tree.Kind.CALL_EXPR); + } + @Test public void tuples() { PyTupleTree empty = parseTuple("()"); @@ -1674,8 +1687,8 @@ public void set_literal() { @Test public void set_comprehension() { setRootRule(PythonGrammar.TEST); - PyListOrSetCompExpressionTree comprehension = - (PyListOrSetCompExpressionTree) parse("{x-1 for x in [42, 43]}", treeMaker::expression); + PyComprehensionExpressionTree comprehension = + (PyComprehensionExpressionTree) parse("{x-1 for x in [42, 43]}", treeMaker::expression); assertThat(comprehension.getKind()).isEqualTo(Tree.Kind.SET_COMPREHENSION); assertThat(comprehension.resultExpression().getKind()).isEqualTo(Tree.Kind.MINUS); assertThat(comprehension.children()).hasSize(2);