PY-10719 Fix add import above import with noinspection comment

GitOrigin-RevId: 6dd79b10e56b9e9e2eee76c83d5974dee43446a1

python/python-psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java

@@ -5,6 +5,7 @@ import com.intellij.lang.ASTNode;
 import com.intellij.lang.injection.InjectedLanguageManager;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.module.Module;
+import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.io.FileUtil;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
@@ -373,21 +374,72 @@ public final class PyPsiUtils {
 
   @NotNull
   public static List<PsiComment> getPrecedingComments(@NotNull PsiElement element, boolean stopAtBlankLine) {
+    return getPrecedingCommentsAndAnchor(element, stopAtBlankLine, true).getFirst();
+  }
+
+  @NotNull
+  private static Pair<List<PsiComment>, PsiElement> getPrecedingCommentsAndAnchor(PsiElement element, boolean stopAtBlankLine,
+                                                                                  boolean strict) {
     final ArrayList<PsiComment> result = new ArrayList<>();
+    PsiElement cursor = element instanceof PsiComment && !strict ? element : element.getPrevSibling();
     while (true) {
       int newLinesCount = 0;
-      for (element = element.getPrevSibling(); element instanceof PsiWhiteSpace; element = element.getPrevSibling()) {
-        newLinesCount += StringUtil.getLineBreakCount(element.getText());
+      while (cursor instanceof PsiWhiteSpace) {
+        newLinesCount += StringUtil.getLineBreakCount(cursor.getText());
+        cursor = cursor.getPrevSibling();
       }
-      if ((stopAtBlankLine && newLinesCount > 1) || !(element instanceof PsiComment)) {
+      if ((stopAtBlankLine && newLinesCount > 1) || !(cursor instanceof PsiComment)) {
         break;
       }
       else {
-        result.add((PsiComment)element);
+        result.add((PsiComment)cursor);
       }
+      cursor = cursor.getPrevSibling();
     }
     Collections.reverse(result);
-    return result;
+    return Pair.create(result, cursor);
+  }
+
+  /**
+   * Return blank-line-separated blocks of consecutive comments preceding the given element.
+   * <p>
+   * For instance, for the following fragment, it will return two blocks of one and two comments.
+   *
+   * <pre>{@code
+   * # comment
+   *
+   * # comment
+   * # comment
+   * def func():
+   *     pass
+   * }</pre>
+   *
+   * Note that in the following case it will additionally return an empty list of comments as the last element
+   * to distinguish between the cases when there is a blank line above the provided element and when there is not.
+   *
+   * <pre>{@code
+   * # comment
+   *
+   * def func():
+   *     pass
+   * }</pre>
+   *
+   */
+  @NotNull
+  public static List<List<PsiComment>> getPrecedingCommentBlocks(@NotNull PsiElement element) {
+    List<List<PsiComment>> blocks = new ArrayList<>();
+    PsiElement anchor = element;
+    do {
+      Pair<List<PsiComment>, PsiElement> blockAndAnchor = getPrecedingCommentsAndAnchor(anchor, true, false);
+      anchor = blockAndAnchor.getSecond();
+      List<PsiComment> block = blockAndAnchor.getFirst();
+      if (block.size() != 0 || anchor instanceof PsiComment) {
+        blocks.add(block);
+      }
+    }
+    while (anchor instanceof PsiComment);
+    Collections.reverse(blocks);
+    return blocks;
   }
 
   @NotNull
@@ -617,9 +669,9 @@ public final class PyPsiUtils {
     final QualifiedName sourceQName = QualifiedName.fromDottedString(source);
 
     return Stream.concat(
-      file.getFromImports().stream().map(PyFromImportStatement::getImportSourceQName),
-      file.getImportTargets().stream().map(PyImportElement::getImportedQName)
-    )
+        file.getFromImports().stream().map(PyFromImportStatement::getImportSourceQName),
+        file.getImportTargets().stream().map(PyImportElement::getImportedQName)
+      )
       .filter(Objects::nonNull)
       .anyMatch(name -> name.matchesPrefix(sourceQName));
   }

python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java

@@ -262,6 +262,14 @@ public final class AddImportHelper {
     if (newImport != null && (priorityAbove == null || priorityAbove.compareTo(priority) < 0)) {
       newImport.putCopyableUserData(PythonCodeStyleService.IMPORT_GROUP_BEGIN, true);
     }
+
+    if (feeler != null) {
+      var anchorComment = getTopmostBoundComment(feeler);
+      if (anchorComment != null) {
+        seeker = anchorComment;
+      }
+    }
+
     if (priorityBelow != null) {
       // actually not necessary because existing import with higher priority (i.e. lower import group)
       // probably should have IMPORT_GROUP_BEGIN flag already, but we add it anyway just for safety
@@ -275,6 +283,26 @@ public final class AddImportHelper {
     return seeker;
   }
 
+  @Nullable
+  private static PsiComment getTopmostBoundComment(@NotNull PsiElement element) {
+    List<List<PsiComment>> commentBlocks = PyPsiUtils.getPrecedingCommentBlocks(element);
+    if (commentBlocks.isEmpty()) return null;
+
+    List<PsiComment> firstBlock = commentBlocks.get(0);
+    PsiComment firstComment = firstBlock.get(0);
+    if (firstComment.getPrevSibling() != null) {
+      return firstComment;
+    }
+
+    PsiComment lastCommentFirstBlock = firstBlock.get(firstBlock.size() - 1);
+    if (PyUtil.isNoinspectionComment(lastCommentFirstBlock)) {
+      return lastCommentFirstBlock;
+    }
+
+    if (commentBlocks.size() == 1) return null;
+    return ContainerUtil.getFirstItem(commentBlocks.get(1));
+  }
+
   private static boolean shouldInsertBefore(@Nullable PyImportStatementBase newImport,
                                             @NotNull PyImportStatementBase existingImport,
                                             @NotNull ImportPriority priority) {

python/python-psi-impl/src/com/jetbrains/python/psi/PyUtil.java

@@ -5,6 +5,7 @@ import com.google.common.collect.Maps;
 import com.intellij.codeInsight.completion.PrioritizedLookupElement;
 import com.intellij.codeInsight.lookup.LookupElement;
 import com.intellij.codeInsight.lookup.LookupElementBuilder;
+import com.intellij.codeInspection.SuppressionUtil;
 import com.intellij.lang.ASTFactory;
 import com.intellij.lang.ASTNode;
 import com.intellij.model.ModelBranch;
@@ -55,6 +56,7 @@ import org.jetbrains.annotations.*;
 import javax.swing.*;
 import java.io.File;
 import java.util.*;
+import java.util.regex.Pattern;
 
 import static com.jetbrains.python.psi.PyFunction.Modifier.CLASSMETHOD;
 import static com.jetbrains.python.psi.PyFunction.Modifier.STATICMETHOD;
@@ -1718,6 +1720,11 @@ public final class PyUtil {
     return directory.findFile(PyNames.INIT_DOT_PY) != null;
   }
 
+  public static boolean isNoinspectionComment(@NotNull PsiComment comment) {
+    Pattern suppressPattern = Pattern.compile(SuppressionUtil.COMMON_SUPPRESS_REGEXP);
+    return suppressPattern.matcher(comment.getText()).find();
+  }
+
   /**
    * This helper class allows to collect various information about AST nodes composing {@link PyStringLiteralExpression}.
    */

python/testData/quickFixes/PyAddImportQuickFixTest/beforeImportAboveBoundCommentsBelowFileCommentBlock/main.py

@@ -0,0 +1,10 @@
+# comment 1
+
+# comment 2
+
+# comment 3
+# comment 4
+import b
+
+print(b)
+<error descr="Unresolved reference 'a'"><caret>a</error>

python/testData/quickFixes/PyAddImportQuickFixTest/beforeImportAboveBoundCommentsBelowFileCommentBlock/main_after.py

@@ -0,0 +1,10 @@
+# comment 1
+import a
+# comment 2
+
+# comment 3
+# comment 4
+import b
+
+print(b)
+a

python/testData/quickFixes/PyAddImportQuickFixTest/beforeImportAboveNoInspectionComment/main.py

@@ -0,0 +1,5 @@
+#noinspection PyUnresolvedReferences
+from b import no_so
+
+
+<error descr="Unresolved reference 'a'"><caret>a</error>

python/testData/quickFixes/PyAddImportQuickFixTest/beforeImportAboveNoInspectionComment/main_after.py

@@ -0,0 +1,6 @@
+import a
+#noinspection PyUnresolvedReferences
+from b import no_so
+
+
+a

python/testData/quickFixes/PyAddImportQuickFixTest/beforeImportBelowFileCommentBlock/main.py

@@ -0,0 +1,7 @@
+#file related comment
+#second file related comment
+import b
+
+
+print(b)
+<error descr="Unresolved reference 'a'"><caret>a</error>

python/testData/quickFixes/PyAddImportQuickFixTest/beforeImportBelowFileCommentBlock/main_after.py

@@ -0,0 +1,8 @@
+#file related comment
+#second file related comment
+import a
+import b
+
+
+print(b)
+a

python/testData/quickFixes/PyAddImportQuickFixTest/beforeImportBelowFileCommentBlockExceptNoInspectionComment/main.py

@@ -0,0 +1,6 @@
+# comment 1
+# noinspection PyUnresolvedReferences
+import b
+
+print(b)
+<error descr="Unresolved reference 'a'"><caret>a</error>

python/testData/quickFixes/PyAddImportQuickFixTest/beforeImportBelowFileCommentBlockExceptNoInspectionComment/main_after.py

@@ -0,0 +1,7 @@
+# comment 1
+import a
+# noinspection PyUnresolvedReferences
+import b
+
+print(b)
+a

python/testData/quickFixes/PyAddImportQuickFixTest/beforeStatementAboveBoundComments/main.py

@@ -0,0 +1,5 @@
+
+# some comment #
+################
+
+print(<error descr="Unresolved reference 'a'"><caret>a</error>)

python/testData/quickFixes/PyAddImportQuickFixTest/beforeStatementAboveBoundComments/main_after.py

@@ -0,0 +1,6 @@
+import a
+
+# some comment #
+################
+
+print(a)

python/testData/quickFixes/PyAddImportQuickFixTest/beforeStatementAboveBoundCommentsBelowFileCommentBlock/main.py

@@ -0,0 +1,7 @@
+# comment 1
+
+# comment 2
+
+pass
+
+<error descr="Unresolved reference 'a'"><caret>a</error>

python/testData/quickFixes/PyAddImportQuickFixTest/beforeStatementAboveBoundCommentsBelowFileCommentBlock/main_after.py

@@ -0,0 +1,8 @@
+# comment 1
+import a
+
+# comment 2
+
+pass
+
+a

python/testData/quickFixes/PyAddImportQuickFixTest/beforeStatementAboveNoInspectionComment/main.py

@@ -0,0 +1,3 @@
+#noinspection PyUnresolvedReferences
+print(no_so)
+<error descr="Unresolved reference 'a'"><caret>a</error>

python/testData/quickFixes/PyAddImportQuickFixTest/beforeStatementAboveNoInspectionComment/main_after.py

@@ -0,0 +1,5 @@
+import a
+
+#noinspection PyUnresolvedReferences
+print(no_so)
+a

python/testData/quickFixes/PyAddImportQuickFixTest/beforeStatementBelowFileCommentBlock/main.py

@@ -0,0 +1,4 @@
+# some comment #
+################
+
+print(<error descr="Unresolved reference 'a'"><caret>a</error>)

python/testData/quickFixes/PyAddImportQuickFixTest/beforeStatementBelowFileCommentBlock/main_after.py

@@ -0,0 +1,5 @@
+# some comment #
+################
+import a
+
+print(a)

python/testData/quickFixes/PyAddImportQuickFixTest/betweenImportsAboveBoundComments/main.py

@@ -0,0 +1,9 @@
+import a
+# comment 1
+
+# comment 2
+# comment 3
+import c
+
+print(a, c)
+<error descr="Unresolved reference 'b'"><caret>b</error>

python/testData/quickFixes/PyAddImportQuickFixTest/betweenImportsAboveBoundComments/main_after.py

@@ -0,0 +1,10 @@
+import a
+import b
+# comment 1
+
+# comment 2
+# comment 3
+import c
+
+print(a, c)
+b

python/testData/quickFixes/PyAddImportQuickFixTest/betweenImportsAboveNoInspectionComment/main.py

@@ -0,0 +1,7 @@
+import a
+#noinspection PyUnresolvedReferences
+from c import no_so
+
+
+print(a)
+<error descr="Unresolved reference 'b'"><caret>b</error>

python/testData/quickFixes/PyAddImportQuickFixTest/betweenImportsAboveNoInspectionComment/main_after.py

@@ -0,0 +1,8 @@
+import a
+import b
+#noinspection PyUnresolvedReferences
+from c import no_so
+
+
+print(a)
+b

python/testData/quickFixes/PyAddImportQuickFixTest/betweenStatementAndImportAboveBoundComments/main.py

@@ -0,0 +1,6 @@
+import argparse
+
+# some comment #
+################
+
+print(argparse, <error descr="Unresolved reference 'a'"><caret>a</error>)

python/testData/quickFixes/PyAddImportQuickFixTest/betweenStatementAndImportAboveBoundComments/main_after.py

@@ -0,0 +1,8 @@
+import argparse
+
+import a
+
+# some comment #
+################
+
+print(argparse, a)

python/testSrc/com/jetbrains/python/PyMiscellaneousPsiOperationsTest.java

@@ -15,13 +15,17 @@
  */
 package com.jetbrains.python;
 
+import com.intellij.psi.PsiComment;
 import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.List;
+
 /**
  * @author Mikhail Golubev
  */
@@ -61,6 +65,57 @@ public class PyMiscellaneousPsiOperationsTest extends PyTestCase {
     checkAddingNameInFromImport("from mod import ()", "bar", null, true, "from mod import (bar)");
   }
 
+  public void testPrecedingImportBlock() {
+    List<List<PsiComment>> blocks;
+    blocks = getPrecedingImportBlocks("# comment\n" +
+                                      "\n" +
+                                      "# comment\n" +
+                                      "# comment\n" +
+                                      "def func(): \n" +
+                                      "    pass");
+    assertSize(2, blocks);
+    assertSize(1, blocks.get(0));
+    assertSize(2, blocks.get(1));
+
+    blocks = getPrecedingImportBlocks("# comment\n" +
+                                      "\n" +
+                                      "# comment\n" +
+                                      "# comment\n" +
+                                      "\n" +
+                                      "def func(): \n" +
+                                      "    pass");
+    assertSize(3, blocks);
+    assertSize(1, blocks.get(0));
+    assertSize(2, blocks.get(1));
+    assertSize(0, blocks.get(2));
+
+    blocks = getPrecedingImportBlocks("def func(): \n" +
+                                      "    pass");
+    assertSize(0, blocks);
+
+    blocks = getPrecedingImportBlocks("# comment\n" +
+                                      "x = 42\n" +
+                                      "\n" +
+                                      "def func(): \n" +
+                                      "    pass");
+    assertSize(0, blocks);
+
+    blocks = getPrecedingImportBlocks("# comment\n" +
+                                      "x = 42\n" +
+                                      "\n" +
+                                      "# comment\n" +
+                                      "def func(): \n" +
+                                      "    pass");
+    assertSize(1, blocks);
+    assertSize(1, blocks.get(0));
+  }
+
+  private List<List<PsiComment>> getPrecedingImportBlocks(@NotNull String text) {
+    PyFile file = assertInstanceOf(myFixture.configureByText("a.py", text), PyFile.class);
+    PyFunction func = file.findTopLevelFunction("func");
+    return PyPsiUtils.getPrecedingCommentBlocks(func);
+  }
+
   private void checkAddingNameInFromImport(@NotNull String fromImport,
                                            @NotNull String newName,
                                            @Nullable String anchorName,

python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java

@@ -61,6 +61,61 @@ public class PyAddImportQuickFixTest extends PyQuickFixTestCase {
     return getCodeStyleSettings().getCustomSettings(PyCodeStyleSettings.class);
   }
 
+  // PY-10719
+  public void testBeforeImportAboveNoInspectionComment() {
+    doMultiFileAutoImportTest("Import 'a'");
+  }
+
+  // PY-10719
+  public void testBetweenImportsAboveNoInspectionComment() {
+    doMultiFileAutoImportTest("Import 'b'");
+  }
+
+  // PY-10719
+  public void testBeforeImportBelowFileCommentBlock() {
+    doMultiFileAutoImportTest("Import 'a'");
+  }
+
+  // PY-10719 PY-26016
+  public void testBetweenStatementAndImportAboveBoundComments() {
+    doMultiFileAutoImportTest("Import 'a'");
+  }
+
+  // PY-10719 PY-26016
+  public void testBeforeStatementBelowFileCommentBlock() {
+    doMultiFileAutoImportTest("Import 'a'");
+  }
+
+  // PY-10719 PY-26016
+  public void testBeforeStatementAboveBoundComments() {
+    doMultiFileAutoImportTest("Import 'a'");
+  }
+
+  // PY-10719 PY-26016
+  public void testBeforeImportAboveBoundCommentsBelowFileCommentBlock() {
+    doMultiFileAutoImportTest("Import 'a'");
+  }
+
+  // PY-10719 PY-26016
+  public void testBetweenImportsAboveBoundComments() {
+    doMultiFileAutoImportTest("Import 'b'");
+  }
+
+  // PY-10719 PY-26016
+  public void testBeforeImportBelowFileCommentBlockExceptNoInspectionComment() {
+    doMultiFileAutoImportTest("Import 'a'");
+  }
+
+  // PY-10719 PY-26016
+  public void testBeforeStatementAboveBoundCommentsBelowFileCommentBlock() {
+    doMultiFileAutoImportTest("Import 'a'");
+  }
+
+  // PY-10719 PY-53487
+  public void testBeforeStatementAboveNoInspectionComment() {
+    doMultiFileAutoImportTest("Import 'a'");
+  }
+
   // PY-19773
   public void testReexportedName() {
     doMultiFileAutoImportTest("Import 'flask.request'");
@@ -364,6 +419,7 @@ public class PyAddImportQuickFixTest extends PyQuickFixTestCase {
     assertNotNull(djangoViewClass);
     assertTrue(PyUserSkeletonsUtil.isUnderUserSkeletonsDirectory(djangoViewClass.getContainingFile()));
   }
+
   private void dumpSdkRootsFileSystemAndIndexResults() {
     dumpSdkRoots();
     VirtualFile skeletonsDir = PyUserSkeletonsUtil.getUserSkeletonsDirectory();