From 3d8d7cda5d0ce46d68ac6a159212ac210e70cabf Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Thu, 9 Mar 2017 15:40:51 +0300 Subject: [PATCH] PY-19837 Preserve trailing comments in optimize imports When there are a few such comments after "from" imports that are going to be combined into one, the first of the comments is inserted on the same line after the statement, and the remaining ones appear below as extra line comments. --- .../imports/PyImportOptimizer.java | 43 ++++++++++++++----- .../optimizeImports/commentsHandling.after.py | 5 ++- .../optimizeImports/commentsHandling.py | 6 +-- 3 files changed, 38 insertions(+), 16 deletions(-) diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java index 467a3591275c..34da95b348ef 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java @@ -29,8 +29,8 @@ 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.intellij.util.containers.*; +import com.intellij.util.containers.HashMap; import com.jetbrains.python.codeInsight.imports.AddImportHelper.ImportPriority; import com.jetbrains.python.formatter.PyCodeStyleSettings; import com.jetbrains.python.inspections.unresolvedReference.PyUnresolvedReferencesInspection; @@ -88,13 +88,15 @@ public class PyImportOptimizer implements ImportOptimizer { private final PyCodeStyleSettings myPySettings; private final List myImportBlock; private final Map> myGroups; - private final MultiMap myImportToComments; + private final MultiMap myImportToLineComments; + private final MultiMap myTransformedImportToTrailingComments; private ImportSorter(@NotNull PyFile file) { myFile = file; myPySettings = CodeStyleSettingsManager.getSettings(myFile.getProject()).getCustomSettings(PyCodeStyleSettings.class); myImportBlock = myFile.getImportBlock(); - myImportToComments = MultiMap.create(); + myImportToLineComments = MultiMap.create(); + myTransformedImportToTrailingComments = MultiMap.create(); myGroups = new EnumMap<>(ImportPriority.class); for (ImportPriority priority : ImportPriority.values()) { myGroups.put(priority, new ArrayList<>()); @@ -136,6 +138,7 @@ public class PyImportOptimizer implements ImportOptimizer { final MultiMap fromImportSources = MultiMap.create(); // Preserve line comments if any final MultiMap precedingComments = MultiMap.create(); + final Map trailingComments = new HashMap<>(); for (PyImportStatementBase statement : imports) { final PyFromImportStatement fromImport = as(statement, PyFromImportStatement.class); @@ -143,6 +146,7 @@ public class PyImportOptimizer implements ImportOptimizer { fromImportSources.putValue(getNormalizedFromImportSource(fromImport), fromImport); } precedingComments.putValues(statement, collectPrecedingLineComments(statement)); + ContainerUtil.putIfNotNull(statement, as(statement.getLastChild(), PsiComment.class), trailingComments); } for (PyImportStatementBase statement : imports) { @@ -159,11 +163,15 @@ public class PyImportOptimizer implements ImportOptimizer { else { topmostImport = newImports.get(0); } - myImportToComments.putValues(topmostImport, precedingComments.get(statement)); + myImportToLineComments.putValues(topmostImport, precedingComments.get(statement)); + final PsiComment trailingComment = trailingComments.get(statement); + if (trailingComment != null) { + myTransformedImportToTrailingComments.putValue(topmostImport, trailingComment); + } result.addAll(newImports); } else { - myImportToComments.putValues(statement, precedingComments.get(statement)); + myImportToLineComments.putValues(statement, precedingComments.get(statement)); result.add(importStatement); } } @@ -202,12 +210,16 @@ public class PyImportOptimizer implements ImportOptimizer { final String importedNames = StringUtil.join(newStatementElements, PsiElement::getText, ", "); final PyFromImportStatement combinedImport = generator.createFromImportStatement(langLevel, source, importedNames, null); ContainerUtil.map2LinkedSet(newStatementElements, e -> (PyImportStatementBase)e.getParent()).forEach(affected -> { - myImportToComments.putValues(combinedImport, precedingComments.get(affected)); + myImportToLineComments.putValues(combinedImport, precedingComments.get(affected)); + final PsiComment trailingComment = trailingComments.get(affected); + if (trailingComment != null) { + myTransformedImportToTrailingComments.putValue(combinedImport, trailingComment); + } }); result.add(combinedImport); } else { - myImportToComments.putValues(fromImport, precedingComments.get(fromImport)); + myImportToLineComments.putValues(fromImport, precedingComments.get(fromImport)); result.add(fromImport); } } @@ -278,11 +290,20 @@ public class PyImportOptimizer implements ImportOptimizer { content.append("\n"); } for (PyImportStatementBase statement : imports) { - //StringUtil.join(myImportToComments.get(statement), PsiElement::getText, "\n", content); - for (PsiComment comment : myImportToComments.get(statement)) { + for (PsiComment comment : myImportToLineComments.get(statement)) { content.append(comment.getText()).append("\n"); } - content.append(statement.getText()).append("\n"); + content.append(statement.getText()); + final Collection trailingComments = myTransformedImportToTrailingComments.get(statement); + if (!trailingComments.isEmpty()) { + content.append(" "); + for (PsiComment comment : trailingComments) { + content.append(comment.getText()).append("\n"); + } + } + else { + content.append("\n"); + } } } diff --git a/python/testData/optimizeImports/commentsHandling.after.py b/python/testData/optimizeImports/commentsHandling.after.py index c6685a8a6eeb..fb51da58822e 100644 --- a/python/testData/optimizeImports/commentsHandling.after.py +++ b/python/testData/optimizeImports/commentsHandling.after.py @@ -4,11 +4,12 @@ import a # trailing comment for normal import # comment for b import b # comment for c, d -import c +import c # trailing comment for c, d import d # comment for name1 and name3 # comment for name2 -from mod import name1, name2, name3 +from mod import name1, name2, name3 # trailing comment for name1, name3 +# trailing comment for name2 # comment for star import from mod2 import * diff --git a/python/testData/optimizeImports/commentsHandling.py b/python/testData/optimizeImports/commentsHandling.py index fd92a238d680..7fa82810804f 100644 --- a/python/testData/optimizeImports/commentsHandling.py +++ b/python/testData/optimizeImports/commentsHandling.py @@ -6,12 +6,12 @@ import b import a # trailing comment for normal import # comment for c, d -import d, c +import d, c # trailing comment for c, d # comment for name2 -from mod import name2 +from mod import name2 # trailing comment for name2 # comment for name1 and name3 -from mod import name1, name3 # trailing comment for "from" import +from mod import name1, name3 # trailing comment for name1, name3 # comment for star import from mod2 import *