diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java index 75d3e40aed81..07d921bd2f60 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java +++ b/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()); + }); + } } } diff --git a/python/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java index 7239f01cc7f3..d7fac1ed3fc4 100644 --- a/python/src/com/jetbrains/python/psi/PyUtil.java +++ b/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 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); diff --git a/python/testData/optimizeImports/blankLineBetweenDocstringAndFirstImportPreserved.after.py b/python/testData/optimizeImports/blankLineBetweenDocstringAndFirstImportPreserved.after.py new file mode 100644 index 000000000000..64cebe733bee --- /dev/null +++ b/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) \ No newline at end of file diff --git a/python/testData/optimizeImports/blankLineBetweenDocstringAndFirstImportPreserved.py b/python/testData/optimizeImports/blankLineBetweenDocstringAndFirstImportPreserved.py new file mode 100644 index 000000000000..f748d726e656 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/python/testData/optimizeImports/blankLineBetweenEncodingDeclarationAndFirstImportPreserved.after.py b/python/testData/optimizeImports/blankLineBetweenEncodingDeclarationAndFirstImportPreserved.after.py new file mode 100644 index 000000000000..1e0b9391bd2f --- /dev/null +++ b/python/testData/optimizeImports/blankLineBetweenEncodingDeclarationAndFirstImportPreserved.after.py @@ -0,0 +1,6 @@ +# -*- coding: utf-8 +from __future__ import unicode_literals, absolute_import + +import os + +os.listdir('.') diff --git a/python/testData/optimizeImports/blankLineBetweenEncodingDeclarationAndFirstImportPreserved.py b/python/testData/optimizeImports/blankLineBetweenEncodingDeclarationAndFirstImportPreserved.py new file mode 100644 index 000000000000..1e0b9391bd2f --- /dev/null +++ b/python/testData/optimizeImports/blankLineBetweenEncodingDeclarationAndFirstImportPreserved.py @@ -0,0 +1,6 @@ +# -*- coding: utf-8 +from __future__ import unicode_literals, absolute_import + +import os + +os.listdir('.') diff --git a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java b/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java index e8a3caa99b37..b198c7ea09da 100644 --- a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java +++ b/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()));