From e72e9b00f07af521b1779a895d1bd536e6ee1fdd Mon Sep 17 00:00:00 2001
From: Guillaume Dequenne <guillaume.dequenne@sonarsource.com>
Date: Mon, 9 Sep 2019 17:42:54 +0200
Subject: [PATCH] Migrate InitReturnsValueCheck to Strongly typed AST

---
 .../python/checks/InitReturnsValueCheck.java  | 72 ++++++++++---------
 1 file changed, 39 insertions(+), 33 deletions(-)

diff --git a/python-checks/src/main/java/org/sonar/python/checks/InitReturnsValueCheck.java b/python-checks/src/main/java/org/sonar/python/checks/InitReturnsValueCheck.java
index d6cfb5ff27..4ddcb41f46 100644
--- a/python-checks/src/main/java/org/sonar/python/checks/InitReturnsValueCheck.java
+++ b/python-checks/src/main/java/org/sonar/python/checks/InitReturnsValueCheck.java
@@ -19,54 +19,60 @@
  */
 package org.sonar.python.checks;
 
-import com.sonar.sslr.api.AstNode;
-import com.sonar.sslr.api.AstNodeType;
-import java.util.Collections;
+import java.util.ArrayList;
 import java.util.List;
-import java.util.Set;
 import org.sonar.check.Rule;
-import org.sonar.python.PythonCheckAstNode;
-import org.sonar.python.api.PythonGrammar;
-import org.sonar.python.api.PythonKeyword;
+import org.sonar.python.PythonSubscriptionCheck;
+import org.sonar.python.api.tree.PyFunctionDefTree;
+import org.sonar.python.api.tree.PyReturnStatementTree;
+import org.sonar.python.api.tree.PyYieldStatementTree;
+import org.sonar.python.api.tree.Tree;
+import org.sonar.python.tree.BaseTreeVisitor;
 
 @Rule(key = InitReturnsValueCheck.CHECK_KEY)
-public class InitReturnsValueCheck extends PythonCheckAstNode {
+public class InitReturnsValueCheck extends PythonSubscriptionCheck {
 
-  public static final String MESSAGE_RETURN = "Remove this return value.";
-  public static final String MESSAGE_YIELD = "Remove this yield statement.";
+  private static final String MESSAGE_RETURN = "Remove this return value.";
+  private static final String MESSAGE_YIELD = "Remove this yield statement.";
 
   public static final String CHECK_KEY = "S2734";
 
   @Override
-  public Set<AstNodeType> subscribedKinds() {
-    return Collections.singleton(PythonGrammar.FUNCDEF);
+  public void initialize(Context context) {
+    context.registerSyntaxNodeConsumer(Tree.Kind.FUNCDEF, ctx -> {
+      PyFunctionDefTree func = (PyFunctionDefTree) ctx.syntaxNode();
+      if (!"__init__".equals(func.name().name())) {
+        return;
+      }
+      ReturnAndYieldVisitor returnAndYieldVisitor = new ReturnAndYieldVisitor();
+      func.body().accept(returnAndYieldVisitor);
+      for (Tree returnNode : returnAndYieldVisitor.returnNodes) {
+        String message = returnNode.is(Tree.Kind.RETURN_STMT) ? MESSAGE_RETURN : MESSAGE_YIELD;
+        ctx.addIssue(returnNode, message);
+      }
+    });
   }
 
-  @Override
-  public void visitNode(AstNode node) {
-    if (!"__init__".equals(node.getFirstChild(PythonGrammar.FUNCNAME).getTokenValue())){
-      return;
+  private static class ReturnAndYieldVisitor extends BaseTreeVisitor {
+
+    List<Tree> returnNodes = new ArrayList<>();
+
+    @Override
+    public void visitFunctionDef(PyFunctionDefTree pyFunctionDefTree) {
+      // Ignore nested function definitions
     }
-    List<AstNode> returnYieldStatements = node.getDescendants(PythonGrammar.YIELD_STMT, PythonGrammar.RETURN_STMT);
-    for (AstNode returnYieldStatement : returnYieldStatements){
-      if (CheckUtils.insideFunction(returnYieldStatement, node) && !returnReturnNone(returnYieldStatement)){
-        raiseIssue(returnYieldStatement);
+
+    @Override
+    public void visitReturnStatement(PyReturnStatementTree pyReturnStatementTree) {
+      if (pyReturnStatementTree.expressions().isEmpty() || pyReturnStatementTree.expressions().get(0).is(Tree.Kind.NONE)) {
+        return;
       }
+      returnNodes.add(pyReturnStatementTree);
     }
-  }
 
-  private static boolean returnReturnNone(AstNode stmt) {
-    return stmt.is(PythonGrammar.RETURN_STMT)
-        && (stmt.getFirstChild(PythonGrammar.TESTLIST) == null
-        || stmt.getFirstChild(PythonGrammar.TESTLIST).getToken().getValue().equals(PythonKeyword.NONE.getValue()));
-  }
-
-  private void raiseIssue(AstNode node) {
-    String message = MESSAGE_RETURN;
-    if (node.is(PythonGrammar.YIELD_STMT)){
-      message = MESSAGE_YIELD;
+    @Override
+    public void visitYieldStatement(PyYieldStatementTree pyYieldStatementTree) {
+      returnNodes.add(pyYieldStatementTree);
     }
-    addIssue(node, message);
   }
 }
-