Strongly Typed AST: introduce Annotated Assignment and rework Assignment

python-checks/src/main/java/org/sonar/python/checks/hotspots/DebugModeCheck.java

@@ -63,7 +63,7 @@ public void initialize(Context context) {
       PyAssignmentStatementTree assignmentStatementTree = (PyAssignmentStatementTree) ctx.syntaxNode();
       for (PyExpressionListTree lhsExpression : assignmentStatementTree.lhsExpressions()) {
         boolean isDebugProperties = lhsExpression.expressions().stream().anyMatch(DebugModeCheck::isDebugIdentifier);
-        if (isDebugProperties && assignmentStatementTree.assignedValues().stream().anyMatch(DebugModeCheck::isTrueLiteral)) {
+        if (isDebugProperties && isTrueLiteral(assignmentStatementTree.assignedValue())) {
           ctx.addIssue(assignmentStatementTree, MESSAGE);
         }
       }

python-squid/src/main/java/org/sonar/python/api/tree/PyAnnotatedAssignmentTree.java

@@ -0,0 +1,40 @@
+/*
+ * SonarQube Python Plugin
+ * Copyright (C) 2011-2019 SonarSource SA
+ * mailto:info AT sonarsource DOT com
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 3 of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+package org.sonar.python.api.tree;
+
+import com.sonar.sslr.api.Token;
+import javax.annotation.CheckForNull;
+
+public interface PyAnnotatedAssignmentTree extends PyStatementTree {
+  // Using PyExpressionTree because variable can be `identifier | attributeref | subscription | slicing`
+  PyExpressionTree variable();
+
+  // TODO: move to a dedicated tree for typed expression
+  Token colonToken();
+
+  // TODO: move to a dedicated tree for typed expression
+  PyExpressionTree annotation();
+
+  @CheckForNull
+  Token equalToken();
+
+  @CheckForNull
+  PyExpressionTree assignedValue();
+}

python-squid/src/main/java/org/sonar/python/api/tree/PyAssignmentStatementTree.java

@@ -23,7 +23,7 @@
 import java.util.List;
 
 public interface PyAssignmentStatementTree extends PyStatementTree {
-  List<PyExpressionTree> assignedValues();
+  PyExpressionTree assignedValue();
 
   List<Token> equalTokens();
 

python-squid/src/main/java/org/sonar/python/api/tree/PyTreeVisitor.java

@@ -150,4 +150,6 @@ public interface PyTreeVisitor {
   void visitDictCompExpression(PyDictCompExpressionTreeImpl tree);
 
   void visitCompoundAssignment(PyCompoundAssignmentStatementTree pyCompoundAssignmentStatementTree);
+
+  void visitAnnotatedAssignment(PyAnnotatedAssignmentTree pyAnnotatedAssignmentTree);
 }

python-squid/src/main/java/org/sonar/python/api/tree/Tree.java

@@ -49,6 +49,8 @@ enum Kind {
 
     ARG_LIST(PyArgListTree.class),
 
+    ANNOTATED_ASSIGNMENT(PyAnnotatedAssignmentTree.class),
+
     ASSERT_STMT(PyAssertStatementTree.class),
 
     ASSIGNMENT_STMT(PyAssignmentStatementTree.class),

python-squid/src/main/java/org/sonar/python/tree/BaseTreeVisitor.java

@@ -23,6 +23,7 @@
 import javax.annotation.Nullable;
 import org.sonar.python.api.tree.PyAliasedNameTree;
 import org.sonar.python.api.tree.PyArgListTree;
+import org.sonar.python.api.tree.PyAnnotatedAssignmentTree;
 import org.sonar.python.api.tree.PyArgumentTree;
 import org.sonar.python.api.tree.PyAssertStatementTree;
 import org.sonar.python.api.tree.PyAssignmentStatementTree;
@@ -318,7 +319,7 @@ public void visitArgument(PyArgumentTree pyArgumentTree) {
   @Override
   public void visitAssignmentStatement(PyAssignmentStatementTree pyAssignmentStatementTree) {
     scan(pyAssignmentStatementTree.lhsExpressions());
-    scan(pyAssignmentStatementTree.assignedValues());
+    scan(pyAssignmentStatementTree.assignedValue());
   }
 
   @Override
@@ -472,4 +473,11 @@ public void visitCompoundAssignment(PyCompoundAssignmentStatementTree pyCompound
     scan(pyCompoundAssignmentStatementTree.lhsExpression());
     scan(pyCompoundAssignmentStatementTree.rhsExpression());
   }
+
+  @Override
+  public void visitAnnotatedAssignment(PyAnnotatedAssignmentTree pyAnnotatedAssignmentTree) {
+    scan(pyAnnotatedAssignmentTree.variable());
+    scan(pyAnnotatedAssignmentTree.annotation());
+    scan(pyAnnotatedAssignmentTree.assignedValue());
+  }
 }

python-squid/src/main/java/org/sonar/python/tree/PyAnnotatedAssignmentTreeImpl.java

@@ -0,0 +1,88 @@
+/*
+ * SonarQube Python Plugin
+ * Copyright (C) 2011-2019 SonarSource SA
+ * mailto:info AT sonarsource DOT com
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 3 of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+package org.sonar.python.tree;
+
+import com.sonar.sslr.api.Token;
+import java.util.Arrays;
+import java.util.List;
+import javax.annotation.CheckForNull;
+import org.sonar.python.api.tree.PyAnnotatedAssignmentTree;
+import org.sonar.python.api.tree.PyExpressionTree;
+import org.sonar.python.api.tree.PyTreeVisitor;
+import org.sonar.python.api.tree.Tree;
+
+public class PyAnnotatedAssignmentTreeImpl extends PyTree implements PyAnnotatedAssignmentTree {
+  private final PyExpressionTree variable;
+  private final Token colonToken;
+  private final PyExpressionTree annotation;
+  private final Token equalToken;
+  private final PyExpressionTree assignedValue;
+
+  public PyAnnotatedAssignmentTreeImpl(PyExpressionTree variable, Token colonToken, PyExpressionTree annotation, Token equalToken, PyExpressionTree assignedValue) {
+    super(variable.firstToken(), assignedValue != null ? assignedValue.lastToken() : annotation.lastToken());
+    this.variable = variable;
+    this.colonToken = colonToken;
+    this.annotation = annotation;
+    this.equalToken = equalToken;
+    this.assignedValue = assignedValue;
+  }
+
+  @Override
+  public PyExpressionTree variable() {
+    return variable;
+  }
+
+  @Override
+  public Token colonToken() {
+    return colonToken;
+  }
+
+  @Override
+  public PyExpressionTree annotation() {
+    return annotation;
+  }
+
+  @CheckForNull
+  @Override
+  public Token equalToken() {
+    return equalToken;
+  }
+
+  @CheckForNull
+  @Override
+  public PyExpressionTree assignedValue() {
+    return assignedValue;
+  }
+
+  @Override
+  public void accept(PyTreeVisitor visitor) {
+    visitor.visitAnnotatedAssignment(this);
+  }
+
+  @Override
+  public List<Tree> children() {
+    return Arrays.asList(variable, annotation, assignedValue);
+  }
+
+  @Override
+  public Kind getKind() {
+    return Kind.ANNOTATED_ASSIGNMENT;
+  }
+}

python-squid/src/main/java/org/sonar/python/tree/PyAssignmentStatementTreeImpl.java

@@ -21,6 +21,7 @@
 
 import com.sonar.sslr.api.AstNode;
 import com.sonar.sslr.api.Token;
+import java.util.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -33,18 +34,18 @@
 public class PyAssignmentStatementTreeImpl extends PyTree implements PyAssignmentStatementTree {
   private final List<Token> assignTokens;
   private final List<PyExpressionListTree> lhsExpressions;
-  private final List<PyExpressionTree> assignedValues;
+  private final PyExpressionTree assignedValue;
 
-  public PyAssignmentStatementTreeImpl(AstNode astNode, List<Token> assignTokens, List<PyExpressionListTree> lhsExpressions, List<PyExpressionTree> assignedValue) {
+  public PyAssignmentStatementTreeImpl(AstNode astNode, List<Token> assignTokens, List<PyExpressionListTree> lhsExpressions,PyExpressionTree assignedValue) {
     super(astNode);
     this.assignTokens = assignTokens;
     this.lhsExpressions = lhsExpressions;
-    this.assignedValues = assignedValue;
+    this.assignedValue = assignedValue;
   }
 
   @Override
-  public List<PyExpressionTree> assignedValues() {
-    return assignedValues;
+  public PyExpressionTree assignedValue() {
+    return assignedValue;
   }
 
   @Override
@@ -69,6 +70,6 @@ public Kind getKind() {
 
   @Override
   public List<Tree> children() {
-    return Stream.of(lhsExpressions, assignedValues).flatMap(List::stream).collect(Collectors.toList());
+    return Stream.of(lhsExpressions, Collections.singletonList(assignedValue)).flatMap(List::stream).collect(Collectors.toList());
   }
 }

python-squid/src/main/java/org/sonar/python/tree/PythonTreeMaker.java

@@ -34,6 +34,7 @@
 import org.sonar.python.api.PythonPunctuator;
 import org.sonar.python.api.PythonTokenType;
 import org.sonar.python.api.tree.PyAliasedNameTree;
+import org.sonar.python.api.tree.PyAnnotatedAssignmentTree;
 import org.sonar.python.api.tree.PyArgListTree;
 import org.sonar.python.api.tree.PyArgumentTree;
 import org.sonar.python.api.tree.PyAssertStatementTree;
@@ -160,6 +161,9 @@ PyStatementTree statement(AstNode astNode) {
     if (astNode.is(PythonGrammar.NONLOCAL_STMT)) {
       return nonlocalStatement(astNode);
     }
+    if (astNode.is(PythonGrammar.EXPRESSION_STMT) && astNode.hasDirectChildren(PythonGrammar.ANNASSIGN)) {
+      return annotatedAssignment(astNode);
+    }
     if (astNode.is(PythonGrammar.EXPRESSION_STMT) && astNode.hasDirectChildren(PythonPunctuator.ASSIGN)) {
       return assignment(astNode);
     }
@@ -184,6 +188,21 @@ PyStatementTree statement(AstNode astNode) {
     throw new IllegalStateException("Statement " + astNode.getType() + " not correctly translated to strongly typed AST");
   }
 
+  public PyAnnotatedAssignmentTree annotatedAssignment(AstNode astNode) {
+    AstNode annAssign = astNode.getFirstChild(PythonGrammar.ANNASSIGN);
+    AstNode colonTokenNode = annAssign.getFirstChild(PythonPunctuator.COLON);
+    PyExpressionTree variable = exprListOrTestList(astNode.getFirstChild(PythonGrammar.TESTLIST_STAR_EXPR));
+    PyExpressionTree annotation = expression(annAssign.getFirstChild(PythonGrammar.TEST));
+    AstNode equalTokenNode = annAssign.getFirstChild(PythonPunctuator.ASSIGN);
+    Token equalToken = null;
+    PyExpressionTree assignedValue = null;
+    if (equalTokenNode != null) {
+      equalToken = equalTokenNode.getToken();
+      assignedValue = expression(equalTokenNode.getNextSibling());
+    }
+    return new PyAnnotatedAssignmentTreeImpl(variable, colonTokenNode.getToken(), annotation, equalToken, assignedValue);
+  }
+
   private PyStatementListTree getStatementListFromSuite(AstNode suite) {
     return new PyStatementListTreeImpl(suite, getStatementsFromSuite(suite), suite.getTokens());
   }
@@ -502,7 +521,6 @@ public PyWhileStatementTreeImpl whileStatement(AstNode astNode) {
   }
 
   public PyExpressionStatementTree expressionStatement(AstNode astNode) {
-    // TODO: handle ANNASSIGN, b.sequence(AUGASSIGN, b.firstOf(YIELD_EXPR, TESTLIST)), b.zeroOrMore("=", b.firstOf(YIELD_EXPR, TESTLIST_STAR_EXPR)
     List<PyExpressionTree> expressions = astNode.getFirstChild(PythonGrammar.TESTLIST_STAR_EXPR).getChildren(PythonGrammar.TEST, PythonGrammar.STAR_EXPR).stream()
       .map(this::expression)
       .collect(Collectors.toList());
@@ -517,11 +535,9 @@ public PyAssignmentStatementTree assignment(AstNode astNode) {
       assignTokens.add(assignNode.getToken());
       lhsExpressions.add(expressionList(assignNode.getPreviousSibling()));
     }
-    AstNode assignedValues = assignNodes.get(assignNodes.size() - 1).getNextSibling();
-    List<PyExpressionTree> expressions = assignedValues.getChildren(PythonGrammar.TEST, PythonGrammar.STAR_EXPR).stream()
-      .map(this::expression)
-      .collect(Collectors.toList());
-    return new PyAssignmentStatementTreeImpl(astNode, assignTokens, lhsExpressions, expressions);
+    AstNode assignedValueNode = assignNodes.get(assignNodes.size() - 1).getNextSibling();
+    PyExpressionTree assignedValue = assignedValueNode.is(PythonGrammar.YIELD_EXPR) ? yieldExpression(assignedValueNode) : exprListOrTestList(assignedValueNode);
+    return new PyAssignmentStatementTreeImpl(astNode, assignTokens, lhsExpressions, assignedValue);
   }
 
   public PyCompoundAssignmentStatementTree compoundAssignment(AstNode astNode) {

python-squid/src/test/java/org/sonar/python/tree/BaseTreeVisitorTest.java

@@ -23,6 +23,7 @@
 import java.util.function.Function;
 import org.junit.Test;
 import org.sonar.python.api.PythonGrammar;
+import org.sonar.python.api.tree.PyAnnotatedAssignmentTree;
 import org.sonar.python.api.tree.PyAssertStatementTree;
 import org.sonar.python.api.tree.PyAssignmentStatementTree;
 import org.sonar.python.api.tree.PyAwaitExpressionTree;
@@ -203,6 +204,17 @@ public void assignement_stmt() {
     verify(visitor).visitExpressionList(tree.lhsExpressions().get(0));
   }
 
+  @Test
+  public void annotated_assignment() {
+    setRootRule(PythonGrammar.EXPRESSION_STMT);
+    PyAnnotatedAssignmentTree tree = parse("a : int = b", treeMaker::annotatedAssignment);
+    BaseTreeVisitor visitor = spy(BaseTreeVisitor.class);
+    visitor.visitAnnotatedAssignment(tree);
+    verify(visitor).visitName((PyNameTree) tree.variable());
+    verify(visitor).visitName((PyNameTree) tree.annotation());
+    verify(visitor).visitName((PyNameTree) tree.assignedValue());
+  }
+
   @Test
   public void lambda() {
     setRootRule(PythonGrammar.LAMBDEF);