From aaef12ce91b76fee4ad8332161b6449ceee70e62 Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Tue, 24 Jan 2017 16:11:46 +0300 Subject: [PATCH] PY-22355 Optimize imports honors new code style settings for "from" imports Namely, it automatically surrounds them with parentheses and inserts a trailing comma after the last imported name if necessary, e.g. when several "from" imports with the same source were combined into a single long statement that needs to be wrapped. --- .../imports/PyImportOptimizer.java | 14 ++++++++++++++ .../PyFromImportPostFormatProcessor.java | 19 ++++++++++++------- .../alphabeticalOrder.after.py | 4 ++-- .../disableAlphabeticalOrder.after.py | 4 ++-- ...esesAndTrailingCommaInFromImports.after.py | 6 ++++++ ...arenthesesAndTrailingCommaInFromImports.py | 4 ++++ .../python/PyOptimizeImportsTest.java | 12 ++++++++++++ 7 files changed, 52 insertions(+), 11 deletions(-) create mode 100644 python/testData/optimizeImports/parenthesesAndTrailingCommaInFromImports.after.py create mode 100644 python/testData/optimizeImports/parenthesesAndTrailingCommaInFromImports.py diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java index 78525760ebd2..8c2aa27bfaf1 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java @@ -22,6 +22,7 @@ import com.intellij.openapi.util.Comparing; import com.intellij.openapi.util.text.StringUtil; import com.intellij.psi.PsiElement; import com.intellij.psi.PsiFile; +import com.intellij.psi.codeStyle.CodeStyleManager; import com.intellij.psi.codeStyle.CodeStyleSettingsManager; import com.intellij.util.containers.ContainerUtil; import com.intellij.util.containers.MultiMap; @@ -226,12 +227,25 @@ public class PyImportOptimizer implements ImportOptimizer { myGroups.put(priority, imports); } } + prepareNewImports(); markGroupStarts(); addImports(myImportBlock.get(0)); myFile.deleteChildRange(myImportBlock.get(0), myImportBlock.get(myImportBlock.size() - 1)); } + private void prepareNewImports() { + for (List imports : myGroups.values()) { + for (int i = 0; i < imports.size(); i++) { + final PyImportStatementBase newImport = imports.get(i); + final CodeStyleManager styleManager = CodeStyleManager.getInstance(newImport.getProject()); + // Some of imports were copied as is and they're still present in the original PSI file + final PyImportStatementBase formatted = (PyImportStatementBase)styleManager.reformat(newImport.copy()); + imports.set(i, formatted); + } + } + } + private void markGroupStarts() { for (List group : myGroups.values()) { boolean firstImportInGroup = true; diff --git a/python/src/com/jetbrains/python/formatter/PyFromImportPostFormatProcessor.java b/python/src/com/jetbrains/python/formatter/PyFromImportPostFormatProcessor.java index 1ac16f0bec76..0b1505386aa5 100644 --- a/python/src/com/jetbrains/python/formatter/PyFromImportPostFormatProcessor.java +++ b/python/src/com/jetbrains/python/formatter/PyFromImportPostFormatProcessor.java @@ -16,10 +16,7 @@ package com.jetbrains.python.formatter; import com.intellij.openapi.util.TextRange; -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.CodeStyleSettings; import com.intellij.psi.codeStyle.CommonCodeStyleSettings; @@ -27,6 +24,7 @@ import com.intellij.psi.impl.source.codeStyle.PostFormatProcessor; import com.intellij.psi.impl.source.codeStyle.PostFormatProcessorHelper; import com.intellij.util.containers.ContainerUtil; import com.jetbrains.python.PyTokenTypes; +import com.jetbrains.python.PythonLanguage; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyPsiUtils; import org.jetbrains.annotations.NotNull; @@ -53,6 +51,7 @@ public class PyFromImportPostFormatProcessor implements PostFormatProcessor { private static class Visitor extends PyRecursiveElementVisitor { private final PostFormatProcessorHelper myHelper; private final List myImportStatements = new ArrayList<>(); + private PsiElement myRootElement; public Visitor(@NotNull CommonCodeStyleSettings settings) { myHelper = new PostFormatProcessorHelper(settings); @@ -83,8 +82,10 @@ public class PyFromImportPostFormatProcessor implements PostFormatProcessor { @NotNull public PsiElement processElement(@NotNull PsiElement element) { + // For some reason smart pointers don't work for non-physical (in particular, generated) elements + myRootElement = element; findAndReplaceFromImports(element); - return element; + return myRootElement; } @NotNull @@ -95,11 +96,15 @@ public class PyFromImportPostFormatProcessor implements PostFormatProcessor { } private void findAndReplaceFromImports(@NotNull PsiElement element) { - if (element.getContainingFile() instanceof PyFile) { + // Copied/generated elements are stored in DummyHolder files, not PyFiles + if (element.getLanguage().isKindOf(PythonLanguage.INSTANCE)) { element.accept(this); Collections.reverse(myImportStatements); for (PyFromImportStatement statement : myImportStatements) { - replaceFromImport(statement); + final PyFromImportStatement newStatement = replaceFromImport(statement); + if (myRootElement == statement) { + myRootElement = newStatement; + } } } } diff --git a/python/testData/optimizeImports/alphabeticalOrder.after.py b/python/testData/optimizeImports/alphabeticalOrder.after.py index b06b7ad267a8..36f8a922ffe8 100644 --- a/python/testData/optimizeImports/alphabeticalOrder.after.py +++ b/python/testData/optimizeImports/alphabeticalOrder.after.py @@ -4,10 +4,10 @@ from __future__ import unicode_literals import sys from datetime import timedelta -import # broken +import # broken import a import b -import foo # broken +import foo # broken import z from from a import C1 diff --git a/python/testData/optimizeImports/disableAlphabeticalOrder.after.py b/python/testData/optimizeImports/disableAlphabeticalOrder.after.py index 25e330235b74..88da7f018d87 100644 --- a/python/testData/optimizeImports/disableAlphabeticalOrder.after.py +++ b/python/testData/optimizeImports/disableAlphabeticalOrder.after.py @@ -11,8 +11,8 @@ from a import C1 from alphabet import D from b import func from -import foo # broken -import # broken +import foo # broken +import # broken from alphabet import * from alphabet import C from alphabet import B, A diff --git a/python/testData/optimizeImports/parenthesesAndTrailingCommaInFromImports.after.py b/python/testData/optimizeImports/parenthesesAndTrailingCommaInFromImports.after.py new file mode 100644 index 000000000000..9e4ec2c836a8 --- /dev/null +++ b/python/testData/optimizeImports/parenthesesAndTrailingCommaInFromImports.after.py @@ -0,0 +1,6 @@ +from statistics import ( + median, + variance, +) + +print( median,variance)#poorly formatted, but shouldn't be affected \ No newline at end of file diff --git a/python/testData/optimizeImports/parenthesesAndTrailingCommaInFromImports.py b/python/testData/optimizeImports/parenthesesAndTrailingCommaInFromImports.py new file mode 100644 index 000000000000..dad4583476e4 --- /dev/null +++ b/python/testData/optimizeImports/parenthesesAndTrailingCommaInFromImports.py @@ -0,0 +1,4 @@ +from statistics import median +from statistics import variance + +print( median,variance)#poorly formatted, but shouldn't be affected \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java b/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java index 340d5ea3df3d..6124cfce4738 100644 --- a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java +++ b/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java @@ -22,6 +22,7 @@ import com.intellij.openapi.projectRoots.Sdk; import com.intellij.openapi.projectRoots.SdkModificator; import com.intellij.openapi.roots.OrderRootType; import com.intellij.openapi.vfs.VirtualFile; +import com.intellij.psi.codeStyle.CommonCodeStyleSettings; import com.jetbrains.python.fixtures.PyTestCase; import com.jetbrains.python.psi.LanguageLevel; import com.jetbrains.python.psi.PyImportStatementBase; @@ -230,6 +231,17 @@ public class PyOptimizeImportsTest extends PyTestCase { runWithLanguageLevel(LanguageLevel.PYTHON36, this::doTest); } + // PY-22355 + public void testParenthesesAndTrailingCommaInFromImports() { + getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_JOIN_FROM_IMPORTS_WITH_SAME_SOURCE = true; + getPythonCodeStyleSettings().FROM_IMPORT_WRAPPING = CommonCodeStyleSettings.WRAP_ALWAYS; + getPythonCodeStyleSettings().FROM_IMPORT_PARENTHESES_FORCE_IF_MULTILINE = true; + getPythonCodeStyleSettings().FROM_IMPORT_NEW_LINE_AFTER_LEFT_PARENTHESIS = true; + getPythonCodeStyleSettings().FROM_IMPORT_NEW_LINE_BEFORE_RIGHT_PARENTHESIS = true; + getPythonCodeStyleSettings().FROM_IMPORT_TRAILING_COMMA_IF_MULTILINE = true; + doTest(); + } + private void doTest() { myFixture.configureByFile(getTestName(true) + ".py"); OptimizeImportsAction.actionPerformedImpl(DataManager.getInstance().getDataContext(myFixture.getEditor().getContentComponent()));