diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java index 34da95b348ef..c9710087fea1 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java @@ -28,9 +28,9 @@ 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.*; +import com.intellij.util.containers.ContainerUtil; import com.intellij.util.containers.HashMap; +import com.intellij.util.containers.MultiMap; import com.jetbrains.python.codeInsight.imports.AddImportHelper.ImportPriority; import com.jetbrains.python.formatter.PyCodeStyleSettings; import com.jetbrains.python.inspections.unresolvedReference.PyUnresolvedReferencesInspection; @@ -145,7 +145,9 @@ public class PyImportOptimizer implements ImportOptimizer { if (fromImport != null && !fromImport.isStarImport()) { fromImportSources.putValue(getNormalizedFromImportSource(fromImport), fromImport); } - precedingComments.putValues(statement, collectPrecedingLineComments(statement)); + if (statement != myImportBlock.get(0)) { + precedingComments.putValues(statement, collectPrecedingLineComments(statement)); + } ContainerUtil.putIfNotNull(statement, as(statement.getLastChild(), PsiComment.class), trailingComments); } @@ -231,7 +233,7 @@ public class PyImportOptimizer implements ImportOptimizer { private static List collectPrecedingLineComments(@NotNull PyImportStatementBase statement) { final List result = new ArrayList<>(); PsiElement prev = PyPsiUtils.getPrevNonWhitespaceSibling(statement); - while ((prev instanceof PsiComment) && onItsOwnLine(prev) && !isShebangComment(((PsiComment)prev))) { + while (prev instanceof PsiComment && isFirstOnLine(prev)) { result.add((PsiComment)prev); prev = PyPsiUtils.getPrevNonWhitespaceSibling(prev); } @@ -239,11 +241,7 @@ public class PyImportOptimizer implements ImportOptimizer { return result; } - private static boolean isShebangComment(@NotNull PsiComment comment) { - return comment.getTextRange().getStartOffset() == 0 && comment.getText().startsWith("#!"); - } - - private static boolean onItsOwnLine(@NotNull PsiElement element) { + private static boolean isFirstOnLine(@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); @@ -275,10 +273,8 @@ public class PyImportOptimizer implements ImportOptimizer { } } final PyImportStatementBase firstImport = myImportBlock.get(0); - final List comments = collectPrecedingLineComments(firstImport); - final PsiElement topmostAnchor = ObjectUtils.notNull(ContainerUtil.getFirstItem(comments), firstImport); - addImportsBefore(topmostAnchor); - myFile.deleteChildRange(topmostAnchor, ContainerUtil.getLastItem(myImportBlock)); + addImportsBefore(firstImport); + myFile.deleteChildRange(firstImport, ContainerUtil.getLastItem(myImportBlock)); } private void addImportsBefore(@NotNull PsiElement anchor) { diff --git a/python/testData/optimizeImports/commentsHandling.after.py b/python/testData/optimizeImports/commentsHandling.after.py index fb51da58822e..d81c62eb6ec7 100644 --- a/python/testData/optimizeImports/commentsHandling.after.py +++ b/python/testData/optimizeImports/commentsHandling.after.py @@ -1,7 +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 new file mode 100644 index 000000000000..74c13fb7685f --- /dev/null +++ b/python/testData/optimizeImports/keepLicenseComment.after.py @@ -0,0 +1,20 @@ +#!/usr/bin/python +# Copyright 2000-2017 JetBrains s.r.o. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 + +from lib import a, b + +print(a, b, path) diff --git a/python/testData/optimizeImports/keepLicenseComment.py b/python/testData/optimizeImports/keepLicenseComment.py new file mode 100644 index 000000000000..4db5cb55a6e9 --- /dev/null +++ b/python/testData/optimizeImports/keepLicenseComment.py @@ -0,0 +1,20 @@ +#!/usr/bin/python +# Copyright 2000-2017 JetBrains s.r.o. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +from lib import a +from lib import b +# specific for import from sys +from sys import path + +print(a, b, path) diff --git a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java b/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java index eb9aef2dde9e..1406a2937a6e 100644 --- a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java +++ b/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java @@ -249,6 +249,12 @@ public class PyOptimizeImportsTest extends PyTestCase { doTest(); } + // PY-19837 + public void testKeepLicenseComment() { + getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_JOIN_FROM_IMPORTS_WITH_SAME_SOURCE = true; + doTest(); + } + private void doTest() { myFixture.configureByFile(getTestName(true) + ".py"); OptimizeImportsAction.actionPerformedImpl(DataManager.getInstance().getDataContext(myFixture.getEditor().getContentComponent()));