PY-23578 Optimize Imports preserves a blank line before the first import import

It happened due to the bug in the platform and only when the number of
blank lines after the import statement is greater than that before it.

As a temporary workaround before the relevant fix in the platform is
accepted I delete old imports as text through the underlying document.

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

@@ -22,10 +22,7 @@ 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;
-import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiWhiteSpace;
+import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
 import com.intellij.psi.util.PsiTreeUtil;
@@ -286,7 +283,7 @@ public class PyImportOptimizer implements ImportOptimizer {
       final PyImportStatementBase lastImport = ContainerUtil.getLastItem(myImportBlock);
       assert lastImport != null;
       addImportsAfter(lastImport);
-      myFile.deleteChildRange(firstElementToRemove, PyPsiUtils.getNextNonWhitespaceSibling(lastImport).getPrevSibling());
+      deleteRangeThroughDocument(firstElementToRemove, PyPsiUtils.getNextNonWhitespaceSibling(lastImport).getPrevSibling());
     }
 
     private void addImportsAfter(@NotNull PsiElement anchor) {
@@ -331,5 +328,11 @@ public class PyImportOptimizer implements ImportOptimizer {
 
       myFile.addRangeAfter(reformattedFile.getFirstChild(), reformattedFile.getLastChild(), anchor);
     }
+
+    private static void deleteRangeThroughDocument(@NotNull PsiElement first, @NotNull PsiElement last) {
+      PyUtil.updateDocumentUnblockedAndCommitted(first, document -> {
+        document.deleteString(first.getTextRange().getStartOffset(), last.getTextRange().getEndOffset());
+      });
+    }
   }
 }

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

@@ -910,6 +910,25 @@ public class PyUtil {
     return as(PyPsiUtils.getPrevNonWhitespaceSibling(statementList), PsiComment.class);
   }
 
+  /**
+   * Retrieve the document from {@link PsiDocumentManager} using the anchor PSI element and pass it to the consumer function
+   * first releasing it from pending PSI modifications it with {@link PsiDocumentManager#doPostponedOperationsAndUnblockDocument(Document)}
+   * and then committing in try/finally block, so that subsequent operations over the PSI can be performed.
+   */
+  public static void updateDocumentUnblockedAndCommitted(@NotNull PsiElement anchor, @NotNull Consumer<Document> consumer) {
+    final PsiDocumentManager manager = PsiDocumentManager.getInstance(anchor.getProject());
+    final Document document = manager.getDocument(anchor.getContainingFile());
+    if (document != null) {
+      manager.doPostponedOperationsAndUnblockDocument(document);
+      try {
+        consumer.consume(document);
+      }
+      finally {
+        manager.commitDocument(document);
+      }
+    }
+  }
+
   public static class KnownDecoratorProviderHolder {
     public static PyKnownDecoratorProvider[] KNOWN_DECORATOR_PROVIDERS = Extensions.getExtensions(PyKnownDecoratorProvider.EP_NAME);
 

python/testData/optimizeImports/blankLineBetweenDocstringAndFirstImportPreserved.after.py

@@ -0,0 +1,9 @@
+# pylint: disable=missing-docstring, invalid-name
+
+"""2016 - Day 1 Puzzle Part 2 tests."""
+
+import sys
+
+from mod import solve, in_between, Point
+
+print(solve, in_between, Point, sys)

python/testData/optimizeImports/blankLineBetweenDocstringAndFirstImportPreserved.py

@@ -0,0 +1,10 @@
+# pylint: disable=missing-docstring, invalid-name
+
+"""2016 - Day 1 Puzzle Part 2 tests."""
+
+import sys
+
+
+from mod import solve, in_between, Point
+
+print(solve, in_between, Point, sys)

python/testData/optimizeImports/blankLineBetweenEncodingDeclarationAndFirstImportPreserved.after.py

@@ -0,0 +1,6 @@
+# -*- coding: utf-8
+from __future__ import unicode_literals, absolute_import
+
+import os
+
+os.listdir('.')

python/testData/optimizeImports/blankLineBetweenEncodingDeclarationAndFirstImportPreserved.py

@@ -0,0 +1,6 @@
+# -*- coding: utf-8
+from __future__ import unicode_literals, absolute_import
+
+import os
+
+os.listdir('.')

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

@@ -88,11 +88,7 @@ public class PyOptimizeImportsTest extends PyTestCase {
 
   // PY-16351
   public void testNoExtraBlankLineAfterImportBlock() {
-    final String testName = getTestName(true);
-    myFixture.copyDirectoryToProject(testName, "");
-    myFixture.configureByFile("main.py");
-    OptimizeImportsAction.actionPerformedImpl(DataManager.getInstance().getDataContext(myFixture.getEditor().getContentComponent()));
-    myFixture.checkResultByFile(testName + "/main.after.py");
+    doMultiFileTest();
   }
 
   // PY-18521
@@ -298,6 +294,24 @@ public class PyOptimizeImportsTest extends PyTestCase {
     doTest();
   }
 
+  // PY-23578
+  public void testBlankLineBetweenDocstringAndFirstImportPreserved() {
+    doTest();
+  }
+
+  // PY-23636
+  public void testBlankLineBetweenEncodingDeclarationAndFirstImportPreserved() {
+    doTest();
+  }
+
+  private void doMultiFileTest() {
+    final String testName = getTestName(true);
+    myFixture.copyDirectoryToProject(testName, "");
+    myFixture.configureByFile("main.py");
+    OptimizeImportsAction.actionPerformedImpl(DataManager.getInstance().getDataContext(myFixture.getEditor().getContentComponent()));
+    myFixture.checkResultByFile(testName + "/main.after.py");
+  }
+
   private void doTest() {
     myFixture.configureByFile(getTestName(true) + ".py");
     OptimizeImportsAction.actionPerformedImpl(DataManager.getInstance().getDataContext(myFixture.getEditor().getContentComponent()));