SONARPY-1025 Add quick fix for S5719 (InstanceAndClassMethodsAtLeastOnePositionalCheck) (#1142)

Author: nils-werner-sonarsource

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

@@ -31,16 +31,32 @@
 import org.sonar.plugins.python.api.symbols.Usage;
 import org.sonar.plugins.python.api.tree.ClassDef;
 import org.sonar.plugins.python.api.tree.FunctionDef;
-import org.sonar.plugins.python.api.tree.Name;
 import org.sonar.plugins.python.api.tree.Parameter;
 import org.sonar.plugins.python.api.tree.Tree;
+import org.sonar.python.quickfix.IssueWithQuickFix;
+import org.sonar.python.quickfix.PythonQuickFix;
 import org.sonar.python.tree.TreeUtils;
 
+import static org.sonar.python.quickfix.PythonTextEdit.insertBefore;
+
 @Rule(key="S5719")
 public class InstanceAndClassMethodsAtLeastOnePositionalCheck extends PythonSubscriptionCheck {
 
   private static final List<String> KNOWN_CLASS_METHODS = Arrays.asList("__new__", "__init_subclass__");
 
+  private enum MethodIssueType {
+    CLASS_METHOD("Add a class parameter", "cls"),
+    REGULAR_METHOD("Add a \"self\" or class parameter", "self", "cls");
+
+    private final String message;
+    private final List<String> insertions;
+
+    MethodIssueType(String message, String... insertions) {
+      this.message = message;
+      this.insertions = Arrays.asList(insertions);
+    }
+  }
+
   private static boolean isUsageInClassBody(Usage usage, ClassDef classDef) {
     // We want all usages that are not function declarations and their closes parent is the class definition
     return usage.kind() != Usage.Kind.FUNC_DECLARATION
@@ -65,19 +81,30 @@ private static void handleFunctionDef(SubscriptionContext ctx, ClassDef classDef
 
     List<String> decoratorNames = functionDef.decorators()
       .stream()
-      .map(decorator ->
-        TreeUtils.decoratorNameFromExpression(decorator.expression())
-      ).filter(Objects::nonNull).collect(Collectors.toList());
+      .map(decorator -> TreeUtils.decoratorNameFromExpression(decorator.expression()))
+      .filter(Objects::nonNull).collect(Collectors.toList());
 
     if (decoratorNames.contains("staticmethod")) {
       return;
     }
 
     String name = functionSymbol.name();
     if (KNOWN_CLASS_METHODS.contains(name) || decoratorNames.contains("classmethod")) {
-      ctx.addIssue(functionDef.defKeyword(), functionDef.rightPar(), "Add a class parameter");
+      addIssue(ctx, functionDef, MethodIssueType.CLASS_METHOD);
     } else {
-      ctx.addIssue(functionDef.defKeyword(), functionDef.rightPar(), "Add a \"self\" or class parameter");
+      addIssue(ctx, functionDef, MethodIssueType.REGULAR_METHOD);
+    }
+  }
+
+  private static void addIssue(SubscriptionContext ctx, FunctionDef functionDef, MethodIssueType type) {
+    IssueWithQuickFix issue = (IssueWithQuickFix) ctx.addIssue(functionDef.defKeyword(), functionDef.rightPar(),
+      type.message);
+
+    for (String insertion : type.insertions) {
+      PythonQuickFix quickFix = PythonQuickFix.newQuickFix(String.format("Add '%s' as the first argument.", insertion))
+        .addTextEdit(insertBefore(functionDef.rightPar(), insertion))
+        .build();
+      issue.addQuickFix(quickFix);
     }
   }
 

python-checks/src/test/java/org/sonar/python/checks/InstanceAndClassMethodsAtLeastOnePositionalCheckTest.java

@@ -20,15 +20,41 @@
 package org.sonar.python.checks;
 
 import org.junit.Test;
+import org.sonar.plugins.python.api.PythonCheck;
+import org.sonar.python.checks.quickfix.PythonQuickFixVerifier;
 import org.sonar.python.checks.utils.PythonCheckVerifier;
 
 public class InstanceAndClassMethodsAtLeastOnePositionalCheckTest {
 
+  private final PythonCheck check = new InstanceAndClassMethodsAtLeastOnePositionalCheck();
   @Test
   public void test() {
-    PythonCheckVerifier.verify(
-      "src/test/resources/checks/instanceAndClassMethodAtLeastOnePositional.py",
-      new InstanceAndClassMethodsAtLeastOnePositionalCheck());
+    PythonCheckVerifier.verify("src/test/resources/checks/instanceAndClassMethodAtLeastOnePositional.py", check);
   }
 
+  @Test
+  public void class_method_quickfix() {
+    String codeWithIssue = "class Foo():\n" +
+      "    @classmethod\n" +
+      "    def bar(): pass";
+    String fixedCodeWithClassParameter = "class Foo():\n" +
+      "    @classmethod\n" +
+      "    def bar(cls): pass";
+
+    PythonQuickFixVerifier.verify(check, codeWithIssue, fixedCodeWithClassParameter);
+  }
+
+  @Test
+  public void regular_method_quickfix() {
+    String codeWithIssue = "class Foo():\n" +
+      "    def bar(): pass";
+    String fixedCodeWithSelfParameter = "class Foo():\n" +
+      "    def bar(self): pass";
+    String fixedCodeWithClassParameter = "class Foo():\n" +
+      "    def bar(cls): pass";
+
+    PythonQuickFixVerifier.verify(check, codeWithIssue,
+      fixedCodeWithSelfParameter,
+      fixedCodeWithClassParameter);
+  }
 }