diff --git a/python/psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java b/python/psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java index cc867d2a69e8..d0bbb4950def 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java +++ b/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 getPrecedingComments(@NotNull PsiElement element) { - PsiComment firstComment = null, lastComment = null; - overComments: + @NotNull + public static List getPrecedingComments(@NotNull PsiElement element) { + return getPrecedingComments(element, true); + } + + @NotNull + public static List getPrecedingComments(@NotNull PsiElement element, boolean stopAtBlankLine) { + final ArrayList 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 diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java index f0dca94ee7f9..8201d8fd4d04 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java +++ b/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 myImportBlock; private final Map> myGroups; - private final MultiMap myNewImportToLineComments; + + private final MultiMap myOldImportToLineComments = MultiMap.create(); + private final MultiMap myOldImportToInnerComments = MultiMap.create(); + private final MultiMap myOldFromImportBySources = MultiMap.create(); + + private final MultiMap myNewImportToLineComments = MultiMap.create(); // Contains trailing and nested comments of modified (split and joined) imports - private final MultiMap myNewImportToInnerComments; + private final MultiMap myNewImportToInnerComments = MultiMap.create(); + private final List 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 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 imports) { + for (PyImportStatementBase statement : imports) { + final PyFromImportStatement fromImport = as(statement, PyFromImportStatement.class); + if (fromImport != null && !fromImport.isStarImport()) { + myOldFromImportBySources.putValue(getNormalizedFromImportSource(fromImport), fromImport); + } + final Couple> 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 transformImportStatements(@NotNull List imports) { final List 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 fromImportSources = MultiMap.create(); - // Preserve line comments if any - final MultiMap precedingComments = MultiMap.create(); - final MultiMap 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 newImports = ContainerUtil.map(importElements, e -> generator.createImportStatement(langLevel, e.getText(), null)); + final List 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 newStatementElements = new ArrayList<>(); - + // We can neither sort, nor combine star imports if (!fromImport.isStarImport()) { - final Collection sameSourceImports = fromImportSources.get(source); + final Collection 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 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 collectPrecedingLineComments(@NotNull PyImportStatementBase statement) { - final List 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> collectPrecedingLineComments(@NotNull PyImportStatementBase statement) { + final List 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 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 importOrdering = Ordering.from(AddImportHelper.getSameGroupImportsComparator(myFile.getProject())); + final Ordering 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 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 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); } } } diff --git a/python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java b/python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java index 5400a9852ab1..8e490b596a09 100644 --- a/python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java +++ b/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 comments = PyPsiUtils.getPrecedingComments(insertionAnchor); - result = insertionAnchor.getParent().addBefore(generatedMethod, comments != null ? comments.getFirst() : insertionAnchor); + final List 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() { diff --git a/python/testData/optimizeImports/commentsHandling.after.py b/python/testData/optimizeImports/commentsHandling.after.py index 83ef3e4e26dd..424b4952b0dc 100644 --- a/python/testData/optimizeImports/commentsHandling.after.py +++ b/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 diff --git a/python/testData/optimizeImports/keepLicenseComment.after.py b/python/testData/optimizeImports/keepLicenseComment.after.py index 74c13fb7685f..c7d85b801daa 100644 --- a/python/testData/optimizeImports/keepLicenseComment.after.py +++ b/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) diff --git a/python/testData/optimizeImports/keepLicenseComment.py b/python/testData/optimizeImports/keepLicenseComment.py index 4db5cb55a6e9..165bcf613faa 100644 --- a/python/testData/optimizeImports/keepLicenseComment.py +++ b/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 diff --git a/python/testData/optimizeImports/stackDanglingCommentsAtEnd.after.py b/python/testData/optimizeImports/stackDanglingCommentsAtEnd.after.py new file mode 100644 index 000000000000..1504c94f3255 --- /dev/null +++ b/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) diff --git a/python/testData/optimizeImports/stackDanglingCommentsAtEnd.py b/python/testData/optimizeImports/stackDanglingCommentsAtEnd.py new file mode 100644 index 000000000000..de0b28ca68ca --- /dev/null +++ b/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) diff --git a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java b/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java index 29bbdf93c479..8269a3b7d994 100644 --- a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java +++ b/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");