SONARPY-2992 Create QuickFix for Rule S7500 (#312)

GitOrigin-RevId: 7f7032c229f2604e83f1a31781816482294b0f57

Author: maksim-grebeniuk-sonarsource

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

@@ -17,22 +17,32 @@
 package org.sonar.python.checks;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import org.sonar.check.Rule;
 import org.sonar.plugins.python.api.PythonSubscriptionCheck;
 import org.sonar.plugins.python.api.SubscriptionContext;
+import org.sonar.plugins.python.api.quickfix.PythonQuickFix;
 import org.sonar.plugins.python.api.tree.ComprehensionExpression;
 import org.sonar.plugins.python.api.tree.ComprehensionFor;
 import org.sonar.plugins.python.api.tree.DictCompExpression;
 import org.sonar.plugins.python.api.tree.Name;
 import org.sonar.plugins.python.api.tree.Tree;
 import org.sonar.plugins.python.api.tree.Tuple;
+import org.sonar.python.quickfix.TextEditUtils;
 import org.sonar.python.tree.TreeUtils;
 
 @Rule(key = "S7500")
 public class UnnecessaryComprehensionCheck extends PythonSubscriptionCheck {
 
+  private static final Map<Tree.Kind, String> COMPREHENSION_TO_QUICK_FIX_FORMAT_MAPPING = Map.of(
+    Tree.Kind.GENERATOR_EXPR, "%s",
+    Tree.Kind.LIST_COMPREHENSION, "list(%s)",
+    Tree.Kind.SET_COMPREHENSION, "set(%s)",
+    Tree.Kind.DICT_COMPREHENSION, "dict(%s)"
+  );
+
   @Override
   public void initialize(Context context) {
     context.registerSyntaxNodeConsumer(Tree.Kind.GENERATOR_EXPR, UnnecessaryComprehensionCheck::checkComprehensionExpression);
@@ -54,7 +64,8 @@ private static void checkComprehensionExpression(SubscriptionContext ctx) {
         && loopExpression instanceof Name loopValueName
         && valueName.name().equals(loopValueName.name())
     ) {
-      ctx.addIssue(comprehension, "Replace this comprehension with passing the iterable to the collection constructor call");
+      var issue = ctx.addIssue(comprehension, "Replace this comprehension with passing the iterable to the collection constructor call");
+      createQuickFix(comprehension, comprehension.comprehensionFor()).ifPresent(issue::addQuickFix);
     }
   }
 
@@ -78,7 +89,17 @@ private static void checkDictComprehensionExpression(SubscriptionContext ctx) {
         && keyName.name().equals(loopKeyName.name())
         && valueName.name().equals(loopValueName.name())
     ) {
-      ctx.addIssue(comprehension, "Replace this comprehension with passing the iterable to the dict constructor call");
+      var issue = ctx.addIssue(comprehension, "Replace this comprehension with passing the iterable to the dict constructor call");
+      createQuickFix(comprehension, comprehension.comprehensionFor()).ifPresent(issue::addQuickFix);
     }
   }
+
+  private static Optional<PythonQuickFix> createQuickFix(Tree comprehension, ComprehensionFor comprehensionFor) {
+    return Optional.ofNullable(TreeUtils.treeToString(comprehensionFor.iterable(), false))
+      .flatMap(iterableString -> Optional.ofNullable(COMPREHENSION_TO_QUICK_FIX_FORMAT_MAPPING.getOrDefault(comprehension.getKind(), null))
+        .map(format -> format.formatted(iterableString))
+        .map(replacementText -> TextEditUtils.replace(comprehension, replacementText))
+        .map(textEdit -> PythonQuickFix.newQuickFix("Replace with collection constructor call").addTextEdit(textEdit))
+        .map(PythonQuickFix.Builder::build));
+  }
 }

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

@@ -16,7 +16,12 @@
  */
 package org.sonar.python.checks;
 
+import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.sonar.python.checks.quickfix.PythonQuickFixVerifier;
 import org.sonar.python.checks.utils.PythonCheckVerifier;
 
 class UnnecessaryComprehensionCheckTest {
@@ -26,4 +31,69 @@ void test() {
     PythonCheckVerifier.verify("src/test/resources/checks/unnecessaryComprehension.py", new UnnecessaryComprehensionCheck());
   }
 
+  @ParameterizedTest
+  @MethodSource("quickFixTestCases")
+  void quickFixTest(String before, String after, String expectedMessage) {
+    var check = new UnnecessaryComprehensionCheck();
+    PythonQuickFixVerifier.verify(check, before, after);
+    PythonQuickFixVerifier.verifyQuickFixMessages(check, before, expectedMessage);
+  }
+
+  @ParameterizedTest
+  @MethodSource("noQuickFixTestCases")
+  void noQuickFixTest(String before) {
+    PythonQuickFixVerifier.verifyNoQuickFixes(new UnnecessaryComprehensionCheck(), before);
+  }
+
+  static Stream<Arguments> quickFixTestCases() {
+    return Stream.of(
+      Arguments.of(
+        """
+          [x for x in some_iterable]
+          """,
+        """
+          list(some_iterable)
+          """,
+        "Replace with collection constructor call"),
+      Arguments.of(
+        """
+          {x for x in some_iterable}
+          """,
+        """
+          set(some_iterable)
+          """,
+        "Replace with collection constructor call"),
+      Arguments.of(
+        """
+          iterable_pairs = [('a', 1), ('b', 2)]
+          {k: v for k, v in iterable_pairs}
+          """,
+        """
+          iterable_pairs = [('a', 1), ('b', 2)]
+          dict(iterable_pairs)
+          """,
+        "Replace with collection constructor call"),
+      Arguments.of(
+        """
+          list(x for x in some_iterable)
+          """,
+        """
+          list(some_iterable)
+          """,
+        "Replace with collection constructor call")
+    );
+  }
+
+  static Stream<Arguments> noQuickFixTestCases() {
+    return Stream.of(
+      Arguments.of(
+        """
+          iterable_pairs = [('a', 1), ('b', 2)]
+          dict_comp = {k: v for k, v in [
+                  ('a', 1),
+                  ('b', 2)
+          ]}
+          """)
+    );
+  }
 }