PY-23125 Optimize imports stacks unbound comments at the end of the import block

Optimize imports now differentiate between "bound" and "unbound" comments
interleaving import statements. The former immediately precede
the import after them without any blank lines between them, while all
the rest are considered the latter, "unbound", comments which are
grouped and inserted after the whole import block. It allows to handle
the comments before the first import more accurately, moving comments
like "# noinspection" as expected, yet leaving licenses, shebangs and
encoding declarations in place if they are separated with a blank line.

Additionally, this is almost identical to the way "isort" utility
handles line comments.

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

@@ -20,7 +20,6 @@ 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.Couple;
 import com.intellij.openapi.util.io.FileUtil;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
@@ -39,6 +38,7 @@ import org.jetbrains.annotations.Nullable;
 
 import java.lang.reflect.Array;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;
 
@@ -400,29 +400,28 @@ public class PyPsiUtils {
    * @param element element comments should be adjacent to
    * @return described range or {@code null} if there are no such comments
    */
-  @Nullable
-  public static Couple<PsiComment> getPrecedingComments(@NotNull PsiElement element) {
-    PsiComment firstComment = null, lastComment = null;
-    overComments:
+  @NotNull
+  public static List<PsiComment> getPrecedingComments(@NotNull PsiElement element) {
+    return getPrecedingComments(element, true);
+  }
+
+  @NotNull
+  public static List<PsiComment> getPrecedingComments(@NotNull PsiElement element, boolean stopAtBlankLine) {
+    final ArrayList<PsiComment> result = new ArrayList<>();
     while (true) {
       int newLinesCount = 0;
       for (element = element.getPrevSibling(); element instanceof PsiWhiteSpace; element = element.getPrevSibling()) {
         newLinesCount += StringUtil.getLineBreakCount(element.getText());
-        if (newLinesCount > 1) {
-          break overComments;
-        }
       }
-      if (element instanceof PsiComment) {
-        if (lastComment == null) {
-          lastComment = (PsiComment)element;
-        }
-        firstComment = (PsiComment)element;
-      }
-      else {
+      if ((stopAtBlankLine && newLinesCount > 1) || !(element instanceof PsiComment)) {
         break;
       }
+      else {
+        result.add((PsiComment)element);
+      }
     }
-    return lastComment == null ? null : Couple.of(firstComment, lastComment);
+    Collections.reverse(result);
+    return result;
   }
 
   @NotNull

python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java

@@ -20,6 +20,7 @@ import com.intellij.codeInspection.LocalInspectionToolSession;
 import com.intellij.lang.ImportOptimizer;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Comparing;
+import com.intellij.openapi.util.Couple;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiComment;
 import com.intellij.psi.PsiElement;
@@ -28,6 +29,7 @@ import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.ObjectUtils;
 import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.containers.MultiMap;
 import com.jetbrains.python.codeInsight.imports.AddImportHelper.ImportPriority;
@@ -87,16 +89,20 @@ public class PyImportOptimizer implements ImportOptimizer {
     private final PyCodeStyleSettings myPySettings;
     private final List<PyImportStatementBase> myImportBlock;
     private final Map<ImportPriority, List<PyImportStatementBase>> myGroups;
-    private final MultiMap<PyImportStatementBase, PsiComment> myNewImportToLineComments;
+    
+    private final MultiMap<PyImportStatementBase, PsiComment> myOldImportToLineComments = MultiMap.create();
+    private final MultiMap<PyImportStatementBase, PsiComment> myOldImportToInnerComments = MultiMap.create();
+    private final MultiMap<String, PyFromImportStatement> myOldFromImportBySources = MultiMap.create();
+    
+    private final MultiMap<PyImportStatementBase, PsiComment> myNewImportToLineComments = MultiMap.create();
     // Contains trailing and nested comments of modified (split and joined) imports
-    private final MultiMap<PyImportStatementBase, PsiComment> myNewImportToInnerComments;
+    private final MultiMap<PyImportStatementBase, PsiComment> myNewImportToInnerComments = MultiMap.create();
+    private final List<PsiComment> myDanglingComments = new ArrayList<>();
 
     private ImportSorter(@NotNull PyFile file) {
       myFile = file;
       myPySettings = CodeStyleSettingsManager.getSettings(myFile.getProject()).getCustomSettings(PyCodeStyleSettings.class);
       myImportBlock = myFile.getImportBlock();
-      myNewImportToLineComments = MultiMap.create();
-      myNewImportToInnerComments = MultiMap.create();
       myGroups = new EnumMap<>(ImportPriority.class);
       for (ImportPriority priority : ImportPriority.values()) {
         myGroups.put(priority, new ArrayList<>());
@@ -107,12 +113,14 @@ public class PyImportOptimizer implements ImportOptimizer {
       if (myImportBlock.isEmpty()) {
         return;
       }
-      
+
+      analyzeImports(myImportBlock);
+
       for (PyImportStatementBase importStatement : myImportBlock) {
         final ImportPriority priority = AddImportHelper.getImportPriority(importStatement);
         myGroups.get(priority).add(importStatement);
       }
-      
+
       boolean hasTransformedImports = false;
       for (ImportPriority priority : ImportPriority.values()) {
         final List<PyImportStatementBase> original = myGroups.get(priority);
@@ -120,12 +128,27 @@ public class PyImportOptimizer implements ImportOptimizer {
         hasTransformedImports |= !original.equals(transformed);
         myGroups.put(priority, transformed);
       }
-      
+
       if (hasTransformedImports || needBlankLinesBetweenGroups() || groupsNotSorted()) {
         applyResults();
       }
     }
-    
+
+    private void analyzeImports(@NotNull List<PyImportStatementBase> imports) {
+      for (PyImportStatementBase statement : imports) {
+        final PyFromImportStatement fromImport = as(statement, PyFromImportStatement.class);
+        if (fromImport != null && !fromImport.isStarImport()) {
+          myOldFromImportBySources.putValue(getNormalizedFromImportSource(fromImport), fromImport);
+        }
+        final Couple<List<PsiComment>> boundAndOthers = collectPrecedingLineComments(statement);
+        myOldImportToLineComments.putValues(statement, boundAndOthers.getFirst());
+        if (statement != myImportBlock.get(0)) {
+          myDanglingComments.addAll(boundAndOthers.getSecond());
+        }
+        myOldImportToInnerComments.putValues(statement, PsiTreeUtil.collectElementsOfType(statement, PsiComment.class));
+      }
+    }
+
     @NotNull
     private List<PyImportStatementBase> transformImportStatements(@NotNull List<PyImportStatementBase> imports) {
       final List<PyImportStatementBase> result = new ArrayList<>();
@@ -133,31 +156,15 @@ public class PyImportOptimizer implements ImportOptimizer {
       final Project project = myFile.getProject();
       final PyElementGenerator generator = PyElementGenerator.getInstance(project);
       final LanguageLevel langLevel = LanguageLevel.forElement(myFile);
-      
-      // Used to combine "from" imports with the same sources
-      final MultiMap<String, PyFromImportStatement> fromImportSources = MultiMap.create();
-      // Preserve line comments if any
-      final MultiMap<PyImportStatementBase, PsiComment> precedingComments = MultiMap.create();
-      final MultiMap<PyImportStatementBase, PsiComment> innerComments = MultiMap.create();
-      
-      for (PyImportStatementBase statement : imports) {
-        final PyFromImportStatement fromImport = as(statement, PyFromImportStatement.class);
-        if (fromImport != null && !fromImport.isStarImport()) {
-          fromImportSources.putValue(getNormalizedFromImportSource(fromImport), fromImport);
-        }
-        if (statement != myImportBlock.get(0)) {
-          precedingComments.putValues(statement, collectPrecedingLineComments(statement));
-        }
-        innerComments.putValues(statement, PsiTreeUtil.collectElementsOfType(statement, PsiComment.class));
-      }
-      
+
       for (PyImportStatementBase statement : imports) {
         if (statement instanceof PyImportStatement) {
           final PyImportStatement importStatement = (PyImportStatement)statement;
           final PyImportElement[] importElements = importStatement.getImportElements();
           // Split combined imports like "import foo, bar as b"
           if (importElements.length > 1) {
-            final List<PyImportStatement> newImports = ContainerUtil.map(importElements, e -> generator.createImportStatement(langLevel, e.getText(), null));
+            final List<PyImportStatement> newImports =
+              ContainerUtil.map(importElements, e -> generator.createImportStatement(langLevel, e.getText(), null));
             final PyImportStatement topmostImport;
             if (myPySettings.OPTIMIZE_IMPORTS_SORT_IMPORTS) {
               topmostImport = Collections.min(newImports, AddImportHelper.getSameGroupImportsComparator(project));
@@ -165,12 +172,12 @@ public class PyImportOptimizer implements ImportOptimizer {
             else {
               topmostImport = newImports.get(0);
             }
-            myNewImportToLineComments.putValues(topmostImport, precedingComments.get(statement));
-            myNewImportToInnerComments.putValues(topmostImport, innerComments.get(statement));
+            myNewImportToLineComments.putValues(topmostImport, myOldImportToLineComments.get(statement));
+            myNewImportToInnerComments.putValues(topmostImport, myOldImportToInnerComments.get(statement));
             result.addAll(newImports);
           }
           else {
-            myNewImportToLineComments.putValues(statement, precedingComments.get(statement));
+            myNewImportToLineComments.putValues(statement, myOldImportToLineComments.get(statement));
             result.add(importStatement);
           }
         }
@@ -178,10 +185,10 @@ public class PyImportOptimizer implements ImportOptimizer {
           final PyFromImportStatement fromImport = (PyFromImportStatement)statement;
           final String source = getNormalizedFromImportSource(fromImport);
           final List<PyImportElement> newStatementElements = new ArrayList<>();
-          
+
           // We can neither sort, nor combine star imports
           if (!fromImport.isStarImport()) {
-            final Collection<PyFromImportStatement> sameSourceImports = fromImportSources.get(source);
+            final Collection<PyFromImportStatement> sameSourceImports = myOldFromImportBySources.get(source);
             if (sameSourceImports.isEmpty()) {
               continue;
             }
@@ -192,7 +199,7 @@ public class PyImportOptimizer implements ImportOptimizer {
                 ContainerUtil.addAll(newStatementElements, sameSourceImport.getImportElements());
               }
               // Remember that we have checked imports with this source already 
-              fromImportSources.remove(source);
+              myOldFromImportBySources.remove(source);
             }
             else if (myPySettings.OPTIMIZE_IMPORTS_SORT_NAMES_IN_FROM_IMPORTS) {
               final List<PyImportElement> originalElements = Arrays.asList(fromImport.getImportElements());
@@ -209,20 +216,20 @@ public class PyImportOptimizer implements ImportOptimizer {
             final String importedNames = StringUtil.join(newStatementElements, ImportSorter::getNormalizedImportElementText, ", ");
             final PyFromImportStatement combinedImport = generator.createFromImportStatement(langLevel, source, importedNames, null);
             ContainerUtil.map2LinkedSet(newStatementElements, e -> (PyImportStatementBase)e.getParent()).forEach(affected -> {
-              myNewImportToLineComments.putValues(combinedImport, precedingComments.get(affected));
-              myNewImportToInnerComments.putValues(combinedImport, innerComments.get(affected));
+              myNewImportToLineComments.putValues(combinedImport, myOldImportToLineComments.get(affected));
+              myNewImportToInnerComments.putValues(combinedImport, myOldImportToInnerComments.get(affected));
             });
             result.add(combinedImport);
           }
           else {
-            myNewImportToLineComments.putValues(fromImport, precedingComments.get(fromImport));
+            myNewImportToLineComments.putValues(fromImport, myOldImportToLineComments.get(fromImport));
             result.add(fromImport);
           }
         }
       }
       return result;
     }
-    
+
     @NotNull
     private static String getNormalizedImportElementText(@NotNull PyImportElement element) {
       // Remove comments, line feeds and backslashes
@@ -230,21 +237,21 @@ public class PyImportOptimizer implements ImportOptimizer {
     }
 
     @NotNull
-    private static List<PsiComment> collectPrecedingLineComments(@NotNull PyImportStatementBase statement) {
-      final List<PsiComment> result = new ArrayList<>();
-      PsiElement prev = PyPsiUtils.getPrevNonWhitespaceSibling(statement);
-      while (prev instanceof PsiComment && isFirstOnLine(prev)) {
-        result.add((PsiComment)prev);
-        prev = PyPsiUtils.getPrevNonWhitespaceSibling(prev);
+    private static Couple<List<PsiComment>> collectPrecedingLineComments(@NotNull PyImportStatementBase statement) {
+      final List<PsiComment> boundComments = PyPsiUtils.getPrecedingComments(statement, true);
+      final PsiComment firstComment = ContainerUtil.getFirstItem(boundComments);
+      if (firstComment != null && isFirstInFile(firstComment)) {
+        return Couple.of(Collections.emptyList(), boundComments);
       }
-      Collections.reverse(result);
-      return result;
+
+      final List<PsiComment> remainingComments = PyPsiUtils.getPrecedingComments(ObjectUtils.notNull(firstComment, statement), false);
+      return Couple.of(boundComments, remainingComments);
     }
 
-    private static boolean isFirstOnLine(@NotNull PsiElement element) {
+    private static boolean isFirstInFile(@NotNull PsiElement element) {
       if (element.getTextRange().getStartOffset() == 0) return true;
-      final PsiWhiteSpace sibling = as(PsiTreeUtil.prevLeaf(element), PsiWhiteSpace.class);
-      return sibling != null && (sibling.textContains('\n') || sibling.getTextRange().getStartOffset() == 0);
+      final PsiWhiteSpace prevWhitespace = as(PsiTreeUtil.prevLeaf(element), PsiWhiteSpace.class);
+      return prevWhitespace != null && prevWhitespace.getTextRange().getStartOffset() == 0;
     }
 
     @NotNull
@@ -256,7 +263,8 @@ public class PyImportOptimizer implements ImportOptimizer {
       if (!myPySettings.OPTIMIZE_IMPORTS_SORT_IMPORTS) {
         return false;
       }
-      final Ordering<PyImportStatementBase> importOrdering = Ordering.from(AddImportHelper.getSameGroupImportsComparator(myFile.getProject()));
+      final Ordering<PyImportStatementBase> importOrdering =
+        Ordering.from(AddImportHelper.getSameGroupImportsComparator(myFile.getProject()));
       return ContainerUtil.exists(myGroups.values(), imports -> !importOrdering.isOrdered(imports));
     }
 
@@ -273,15 +281,20 @@ public class PyImportOptimizer implements ImportOptimizer {
         }
       }
       final PyImportStatementBase firstImport = myImportBlock.get(0);
-      addImportsBefore(firstImport);
-      myFile.deleteChildRange(firstImport, ContainerUtil.getLastItem(myImportBlock));
+      final List<PsiComment> boundComments = collectPrecedingLineComments(firstImport).getFirst();
+      final PsiElement insertionAnchor = boundComments.isEmpty() ? firstImport : boundComments.get(0);
+      addImportsBefore(insertionAnchor);
+      // Remove together with extra whitespaces preceding
+      final PsiElement nonWhitespaceBefore = PyPsiUtils.getPrevNonWhitespaceSibling(insertionAnchor);
+      myFile.deleteChildRange(nonWhitespaceBefore != null ? nonWhitespaceBefore.getNextSibling() : insertionAnchor, 
+                              ContainerUtil.getLastItem(myImportBlock));
     }
 
     private void addImportsBefore(@NotNull PsiElement anchor) {
       final StringBuilder content = new StringBuilder();
-      
+
       for (List<PyImportStatementBase> imports : myGroups.values()) {
-        if (content.length() > 0) {
+        if (content.length() > 0 && !imports.isEmpty()) {
           // one extra blank line between import groups according to PEP 8
           content.append("\n");
         }
@@ -302,7 +315,14 @@ public class PyImportOptimizer implements ImportOptimizer {
           }
         }
       }
-      
+
+      if (!myDanglingComments.isEmpty()) {
+        content.append("\n");
+        for (PsiComment comment : myDanglingComments) {
+          content.append(comment.getText()).append("\n");
+        }
+      }
+
       final Project project = anchor.getProject();
       final PyElementGenerator generator = PyElementGenerator.getInstance(project);
       PyFile file = (PyFile)generator.createDummyFile(LanguageLevel.forElement(anchor), content.toString());
@@ -312,7 +332,7 @@ public class PyImportOptimizer implements ImportOptimizer {
 
       final PyImportStatementBase lastImport = ContainerUtil.getLastItem(newImportBlock);
       assert lastImport != null;
-      myFile.addRangeBefore(file.getFirstChild(), lastImport, anchor);
+      myFile.addRangeBefore(file.getFirstChild(), file.getLastChild(), anchor);
     }
   }
 }

python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java

@@ -25,7 +25,6 @@ import com.intellij.openapi.command.CommandProcessor;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.ui.Messages;
-import com.intellij.openapi.util.Couple;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.util.text.StringUtil;
@@ -497,8 +496,8 @@ public class PyExtractMethodUtil {
       final PsiElement target = parent instanceof PyClass ? ((PyClass)parent).getStatementList() : parent;
       final PsiElement insertionAnchor = PyPsiUtils.getParentRightBefore(anchor, target);
       assert insertionAnchor != null;
-      final Couple<PsiComment> comments = PyPsiUtils.getPrecedingComments(insertionAnchor);
-      result = insertionAnchor.getParent().addBefore(generatedMethod, comments != null ? comments.getFirst() : insertionAnchor);
+      final List<PsiComment> comments = PyPsiUtils.getPrecedingComments(insertionAnchor);
+      result = insertionAnchor.getParent().addBefore(generatedMethod, !comments.isEmpty() ? comments.get(0) : insertionAnchor);
     }
     // to ensure correct reformatting, mark the entire method as generated
     result.accept(new PsiRecursiveElementVisitor() {

python/testData/optimizeImports/commentsHandling.after.py

@@ -1,8 +1,8 @@
 #!/usr/bin/python
 
-# comment for b
 # comment for a
 import a  # trailing comment for normal import 
+# comment for b
 import b
 # comment for c, d
 import c  # trailing comment for c, d

python/testData/optimizeImports/keepLicenseComment.after.py

@@ -12,9 +12,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 # specific for import from sys
 from sys import path
 
+# noinspection PyUnresolvedReferences
+import unresolved
 from lib import a, b
 
 print(a, b, path)

python/testData/optimizeImports/keepLicenseComment.py

@@ -12,6 +12,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+# noinspection PyUnresolvedReferences
+import unresolved
 from lib import a
 from lib import b
 # specific for import from sys

python/testData/optimizeImports/stackDanglingCommentsAtEnd.after.py

@@ -0,0 +1,23 @@
+#!/bin/python
+# top-level comment
+# continuation of the top-level comment
+
+# comment after the top-level one
+
+# comment for abc
+import abc
+# comment for sys
+import sys
+
+# comment for a 1
+# comment for a 2
+import a
+# comment for b
+import b
+
+# dangling comment 1
+# dangling comment 2
+# dangling comment 3
+# dangling comment 4
+
+print(a, b, abc, sys)

python/testData/optimizeImports/stackDanglingCommentsAtEnd.py

@@ -0,0 +1,28 @@
+#!/bin/python
+# top-level comment
+# continuation of the top-level comment
+
+# comment after the top-level one
+
+# comment for b
+import b
+
+# dangling comment 1
+
+# comment for a 1
+# comment for a 2
+import a
+
+# dangling comment 2
+
+# comment for sys
+import sys
+
+# dangling comment 3
+
+# dangling comment 4
+
+# comment for abc
+import abc
+
+print(a, b, abc, sys)

python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java

@@ -272,6 +272,10 @@ public class PyOptimizeImportsTest extends PyTestCase {
     getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_JOIN_FROM_IMPORTS_WITH_SAME_SOURCE = true;
     doTest();
   }
+  
+  public void testStackDanglingCommentsAtEnd() {
+    doTest();
+  }
 
   private void doTest() {
     myFixture.configureByFile(getTestName(true) + ".py");