SONARPY-380 Migrate MissingDocstringCheck to strongly typed AST

python-checks/src/main/java/org/sonar/python/checks/MissingDocstringCheck.java

@@ -19,18 +19,21 @@
  */
 package org.sonar.python.checks;
 
-import com.sonar.sslr.api.AstNode;
-import com.sonar.sslr.api.AstNodeType;
 import com.sonar.sslr.api.Token;
-import java.util.Set;
 import java.util.regex.Pattern;
+import javax.annotation.CheckForNull;
 import org.sonar.check.Rule;
-import org.sonar.python.DocstringExtractor;
-import org.sonar.python.PythonCheckAstNode;
-import org.sonar.python.api.PythonGrammar;
+import org.sonar.python.PythonSubscriptionCheck;
+import org.sonar.python.SubscriptionContext;
+import org.sonar.python.api.tree.PyClassDefTree;
+import org.sonar.python.api.tree.PyFileInputTree;
+import org.sonar.python.api.tree.PyFunctionDefTree;
+import org.sonar.python.api.tree.PyNameTree;
+import org.sonar.python.api.tree.Tree;
+import org.sonar.python.api.tree.Tree.Kind;
 
 @Rule(key = MissingDocstringCheck.CHECK_KEY)
-public class MissingDocstringCheck extends PythonCheckAstNode {
+public class MissingDocstringCheck extends PythonSubscriptionCheck {
 
   public static final String CHECK_KEY = "S1720";
 
@@ -52,53 +55,57 @@ private enum DeclarationType {
   }
 
   @Override
-  public Set<AstNodeType> subscribedKinds() {
-    return DocstringExtractor.DOCUMENTABLE_NODE_TYPES;
+  public void initialize(Context context) {
+    context.registerSyntaxNodeConsumer(Kind.FILE_INPUT, ctx -> checkDocString(ctx, ((PyFileInputTree) ctx.syntaxNode()).docstring()));
+    context.registerSyntaxNodeConsumer(Kind.FUNCDEF, ctx -> checkDocString(ctx, ((PyFunctionDefTree) ctx.syntaxNode()).docstring()));
+    context.registerSyntaxNodeConsumer(Kind.CLASSDEF, ctx -> checkDocString(ctx, ((PyClassDefTree) ctx.syntaxNode()).docstring()));
   }
 
-  @Override
-  public void visitNode(AstNode astNode) {
-    DeclarationType type = getType(astNode);
-    Token docstring = DocstringExtractor.extractDocstring(astNode);
+  private void checkDocString(SubscriptionContext ctx, @CheckForNull Token docstring) {
+    Tree tree = ctx.syntaxNode();
+    DeclarationType type = getType(tree);
     if (docstring == null) {
-      raiseIssueNoDocstring(astNode, type);
+      raiseIssueNoDocstring(tree, type, ctx);
     } else if (EMPTY_STRING_REGEXP.matcher(docstring.getValue()).matches()) {
-      raiseIssue(astNode, MESSAGE_EMPTY_DOCSTRING, type);
+      raiseIssue(tree, MESSAGE_EMPTY_DOCSTRING, type, ctx);
     }
   }
 
-  private static DeclarationType getType(AstNode node) {
-    if (node.is(PythonGrammar.FUNCDEF)) {
-      if (CheckUtils.isMethodDefinition(node)) {
+  private static DeclarationType getType(Tree tree) {
+    if (tree.is(Kind.FUNCDEF)) {
+      if (((PyFunctionDefTree) tree).isMethodDefinition()) {
         return DeclarationType.METHOD;
       } else {
         return DeclarationType.FUNCTION;
       }
-    } else if (node.is(PythonGrammar.CLASSDEF)) {
+    } else if (tree.is(Kind.CLASSDEF)) {
       return DeclarationType.CLASS;
     } else {
-      // node is PythonGrammar.FILE_INPUT
+      // tree is FILE_INPUT
       return DeclarationType.MODULE;
     }
   }
 
-  private void raiseIssueNoDocstring(AstNode astNode, DeclarationType type) {
+  private void raiseIssueNoDocstring(Tree tree, DeclarationType type, SubscriptionContext ctx) {
     if (type != DeclarationType.METHOD) {
-      raiseIssue(astNode, MESSAGE_NO_DOCSTRING, type);
+      raiseIssue(tree, MESSAGE_NO_DOCSTRING, type, ctx);
     }
   }
 
-  private void raiseIssue(AstNode astNode, String message, DeclarationType type) {
+  private void raiseIssue(Tree tree, String message, DeclarationType type, SubscriptionContext ctx) {
     String finalMessage = String.format(message, type.value);
     if (type != DeclarationType.MODULE) {
-      addIssue(getNameNode(astNode), finalMessage);
+      ctx.addIssue(getNameNode(tree), finalMessage);
     } else {
-      addFileIssue(finalMessage);
+      ctx.addFileIssue(finalMessage);
     }
   }
 
-  private static AstNode getNameNode(AstNode astNode) {
-    return astNode.getFirstChild(PythonGrammar.FUNCNAME, PythonGrammar.CLASSNAME);
+  private static PyNameTree getNameNode(Tree tree) {
+    if (tree.is(Kind.FUNCDEF)) {
+      return ((PyFunctionDefTree) tree).name();
+    }
+    return ((PyClassDefTree) tree).name();
   }
 
 }

python-squid/src/main/java/org/sonar/python/SubscriptionContext.java

@@ -31,6 +31,8 @@ public interface SubscriptionContext {
 
   PythonCheck.PreciseIssue addIssue(Token token, @Nullable String message);
 
+  PythonCheck.PreciseIssue addFileIssue(String finalMessage);
+
   SymbolTable symbolTable();
 
   PythonFile pythonFile();

python-squid/src/main/java/org/sonar/python/SubscriptionVisitor.java

@@ -27,6 +27,7 @@
 import java.util.List;
 import java.util.function.Consumer;
 import javax.annotation.Nullable;
+import org.sonar.python.api.tree.PyFileInputTree;
 import org.sonar.python.api.tree.Tree;
 import org.sonar.python.api.tree.Tree.Kind;
 import org.sonar.python.semantic.SymbolTable;
@@ -40,8 +41,9 @@ public class SubscriptionVisitor extends BaseTreeVisitor {
 
   public static void analyze(Collection<PythonSubscriptionCheck> checks, PythonVisitorContext pythonVisitorContext) {
     SubscriptionVisitor subscriptionVisitor = new SubscriptionVisitor(checks, pythonVisitorContext);
-    if (pythonVisitorContext.rootTree() != null) {
-      pythonVisitorContext.rootTree().accept(subscriptionVisitor);
+    PyFileInputTree rootTree = pythonVisitorContext.rootTree();
+    if (rootTree != null) {
+      subscriptionVisitor.scan(rootTree);
     }
   }
 
@@ -96,6 +98,13 @@ public PythonCheck.PreciseIssue addIssue(Token token, @Nullable String message)
       return newIssue;
     }
 
+    @Override
+    public PythonCheck.PreciseIssue addFileIssue(String message) {
+      PythonCheck.PreciseIssue newIssue = new PythonCheck.PreciseIssue(check, IssueLocation.atFileLevel(message));
+      pythonVisitorContext.addIssue(newIssue);
+      return newIssue;
+    }
+
     @Override
     public SymbolTable symbolTable() {
       return pythonVisitorContext.symbolTable();

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

@@ -43,4 +43,7 @@ public interface PyClassDefTree extends PyStatementTree {
   Token colon();
 
   PyStatementListTree body();
+
+  @CheckForNull
+  Token docstring();
 }

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

@@ -19,9 +19,13 @@
  */
 package org.sonar.python.api.tree;
 
+import com.sonar.sslr.api.Token;
 import javax.annotation.CheckForNull;
 
 public interface PyFileInputTree extends Tree {
   @CheckForNull
   PyStatementListTree statements();
+
+  @CheckForNull
+  Token docstring();
 }

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

@@ -60,4 +60,7 @@ public interface PyFunctionDefTree extends PyStatementTree {
 
   boolean isMethodDefinition();
 
+  @CheckForNull
+  Token docstring();
+
 }

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

@@ -35,12 +35,14 @@ public class PyClassDefTreeImpl extends PyTree implements PyClassDefTree {
   private final PyNameTree name;
   private final PyArgListTree args;
   private final PyStatementListTree body;
+  private final Token docstring;
 
-  public PyClassDefTreeImpl(AstNode astNode, PyNameTree name, PyArgListTree args, PyStatementListTree body) {
+  public PyClassDefTreeImpl(AstNode astNode, PyNameTree name, PyArgListTree args, PyStatementListTree body, Token docstring) {
     super(astNode);
     this.name = name;
     this.args = args;
     this.body = body;
+    this.docstring = docstring;
   }
 
   @Override
@@ -95,4 +97,10 @@ public Token colon() {
   public PyStatementListTree body() {
     return body;
   }
+
+  @CheckForNull
+  @Override
+  public Token docstring() {
+    return docstring;
+  }
 }

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

@@ -20,6 +20,7 @@
 package org.sonar.python.tree;
 
 import com.sonar.sslr.api.AstNode;
+import com.sonar.sslr.api.Token;
 import javax.annotation.CheckForNull;
 import javax.annotation.Nullable;
 import org.sonar.python.api.tree.PyFileInputTree;
@@ -30,10 +31,12 @@
 public class PyFileInputTreeImpl extends PyTree implements PyFileInputTree {
 
   private final PyStatementListTree statements;
+  private final Token docstring;
 
-  public PyFileInputTreeImpl(AstNode astNode, @Nullable PyStatementListTree statements) {
+  public PyFileInputTreeImpl(AstNode astNode, @Nullable PyStatementListTree statements, Token docstring) {
     super(astNode);
     this.statements = statements;
+    this.docstring = docstring;
   }
 
   @Override
@@ -47,6 +50,12 @@ public PyStatementListTree statements() {
     return statements;
   }
 
+  @CheckForNull
+  @Override
+  public Token docstring() {
+    return docstring;
+  }
+
   @Override
   public void accept(PyTreeVisitor visitor) {
     visitor.visitFileInput(this);

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

@@ -37,13 +37,15 @@ public class PyFunctionDefTreeImpl extends PyTree implements PyFunctionDefTree {
   private final PyTypedArgListTree typedArgs;
   private final PyStatementListTree body;
   private final boolean isMethodDefinition;
+  private final Token docstring;
 
-  public PyFunctionDefTreeImpl(AstNode astNode, PyNameTree name, PyTypedArgListTree typedArgs, PyStatementListTree body, boolean isMethodDefinition) {
+  public PyFunctionDefTreeImpl(AstNode astNode, PyNameTree name, PyTypedArgListTree typedArgs, PyStatementListTree body, boolean isMethodDefinition, Token docstring) {
     super(astNode);
     this.name = name;
     this.typedArgs = typedArgs;
     this.body = body;
     this.isMethodDefinition = isMethodDefinition;
+    this.docstring = docstring;
   }
 
   @Override
@@ -115,6 +117,12 @@ public boolean isMethodDefinition() {
     return isMethodDefinition;
   }
 
+  @CheckForNull
+  @Override
+  public Token docstring() {
+    return docstring;
+  }
+
   @Override
   public Kind getKind() {
     return Kind.FUNCDEF;

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

@@ -26,6 +26,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
+import org.sonar.python.DocstringExtractor;
 import org.sonar.python.api.PythonGrammar;
 import org.sonar.python.api.PythonKeyword;
 import org.sonar.python.api.PythonPunctuator;
@@ -81,7 +82,7 @@ public class PythonTreeMaker {
   public PyFileInputTree fileInput(AstNode astNode) {
     List<PyStatementTree> statements = getStatements(astNode).stream().map(this::statement).collect(Collectors.toList());
     PyStatementListTreeImpl statementList = statements.isEmpty() ? null : new PyStatementListTreeImpl(astNode, statements);
-    return new PyFileInputTreeImpl(astNode, statementList);
+    return new PyFileInputTreeImpl(astNode, statementList, DocstringExtractor.extractDocstring(astNode));
   }
 
   PyStatementTree statement(AstNode astNode) {
@@ -405,7 +406,7 @@ public PyFunctionDefTree funcDefStatement(AstNode astNode) {
     }
 
     PyStatementListTree body = getStatementListFromSuite(astNode.getFirstChild(PythonGrammar.SUITE));
-    return new PyFunctionDefTreeImpl(astNode, name, typedArgs, body, isMethodDefinition(astNode));
+    return new PyFunctionDefTreeImpl(astNode, name, typedArgs, body, isMethodDefinition(astNode), DocstringExtractor.extractDocstring(astNode));
   }
 
   private static boolean isMethodDefinition(AstNode node) {
@@ -424,7 +425,7 @@ public PyClassDefTree classDefStatement(AstNode astNode) {
     // TODO argList
     PyArgListTree args = null;
     PyStatementListTree body = getStatementListFromSuite(astNode.getFirstChild(PythonGrammar.SUITE));
-    return new PyClassDefTreeImpl(astNode, name, args, body);
+    return new PyClassDefTreeImpl(astNode, name, args, body, DocstringExtractor.extractDocstring(astNode));
   }
 
   private PyNameTree name(AstNode astNode) {

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

@@ -82,6 +82,15 @@ public class PythonTreeMakerTest extends RuleTest {
   public void fileInputTreeOnEmptyFile() {
     PyFileInputTree pyTree = parse("", treeMaker::fileInput);
     assertThat(pyTree.statements()).isNull();
+    assertThat(pyTree.docstring()).isNull();
+
+    pyTree = parse("\"\"\"\n" +
+      "This is a module docstring\n" +
+      "\"\"\"", treeMaker::fileInput);
+    assertThat(pyTree.docstring().getValue()).isEqualTo("\"\"\"\n" +
+      "This is a module docstring\n" +
+      "\"\"\"");
+
   }
 
   @Test
@@ -542,6 +551,7 @@ public void funcdef_statement() {
     assertThat(functionDefTree.body().statements().get(0).is(Tree.Kind.PASS_STMT)).isTrue();
     assertThat(functionDefTree.typedArgs()).isNull();
     assertThat(functionDefTree.isMethodDefinition()).isFalse();
+    assertThat(functionDefTree.docstring()).isNull();
     // TODO
     assertThat(functionDefTree.decorators()).isNull();
     assertThat(functionDefTree.asyncKeyword()).isNull();
@@ -563,6 +573,13 @@ public void funcdef_statement() {
 
     functionDefTree = parse("def func(x : int, y): pass", treeMaker::funcDefStatement);
     assertThat(functionDefTree.typedArgs().arguments()).hasSize(2);
+
+    functionDefTree = parse("def func(x : int, y):\n  \"\"\"\n" +
+      "This is a function docstring\n" +
+      "\"\"\"\n  pass", treeMaker::funcDefStatement);
+    assertThat(functionDefTree.docstring().getValue()).isEqualTo("\"\"\"\n" +
+      "This is a function docstring\n" +
+      "\"\"\"");
   }
 
   @Test
@@ -571,6 +588,7 @@ public void classdef_statement() {
     AstNode astNode = p.parse("class clazz: pass");
     PyClassDefTree classDefTree = treeMaker.classDefStatement(astNode);
     assertThat(classDefTree.name()).isNotNull();
+    assertThat(classDefTree.docstring()).isNull();
     assertThat(classDefTree.name().name()).isEqualTo("clazz");
     assertThat(classDefTree.body().statements()).hasSize(1);
     assertThat(classDefTree.body().statements().get(0).is(Tree.Kind.PASS_STMT)).isTrue();
@@ -581,6 +599,13 @@ public void classdef_statement() {
     classDefTree = treeMaker.classDefStatement(astNode);
     PyFunctionDefTree funcDef = (PyFunctionDefTree) classDefTree.body().statements().get(0);
     assertThat(funcDef.isMethodDefinition()).isTrue();
+
+
+    astNode = p.parse("class ClassWithDocstring:\n" +
+      "\t\"\"\"This is a docstring\"\"\"\n" +
+      "\tpass");
+    classDefTree = treeMaker.classDefStatement(astNode);
+    assertThat(classDefTree.docstring().getValue()).isEqualTo("\"\"\"This is a docstring\"\"\"");
   }
 
   @Test