From 6bb86279d1e11b9cc157d0ec597339f8b7e03a8a Mon Sep 17 00:00:00 2001 From: Dmitry Jemerov Date: Tue, 2 Oct 2012 13:52:55 +0200 Subject: [PATCH] optimize imports sorts them according to PEP-8 (PY-2367) --- .../python/psi/PyElementGenerator.java | 4 +- .../src/com/jetbrains/python/psi/PyFile.java | 5 + .../python/PythonFoldingBuilder.java | 40 ++----- .../codeInsight/imports/AddImportHelper.java | 16 +-- .../imports/ImportFromExistingAction.java | 2 +- .../imports/PyImportOptimizer.java | 109 +++++++++++++++++- .../jetbrains/python/formatter/PyBlock.java | 11 ++ .../psi/impl/PyElementGeneratorImpl.java | 9 +- .../jetbrains/python/psi/impl/PyFileImpl.java | 31 +++++ python/testData/folding/importBlock.py | 7 ++ .../testData/optimizeImports/order.after.py | 9 ++ python/testData/optimizeImports/order.py | 7 ++ .../testData/optimizeImports/split.after.py | 5 + python/testData/optimizeImports/split.py | 4 + .../com/jetbrains/python/PyFoldingTest.java | 4 + .../python/PyOptimizeImportsTest.java | 10 +- 16 files changed, 225 insertions(+), 48 deletions(-) create mode 100644 python/testData/folding/importBlock.py create mode 100644 python/testData/optimizeImports/order.after.py create mode 100644 python/testData/optimizeImports/order.py create mode 100644 python/testData/optimizeImports/split.after.py create mode 100644 python/testData/optimizeImports/split.py diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java b/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java index ccbc4972251c..8304e3e691aa 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java +++ b/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java @@ -56,9 +56,9 @@ public abstract class PyElementGenerator { @NotNull public abstract PyCallExpression createCallExpression(final LanguageLevel langLevel, String functionName); - public abstract PyImportStatement createImportStatementFromText(String text); + public abstract PyImportStatement createImportStatementFromText(final LanguageLevel languageLevel, String text); - public abstract PyImportElement createImportElement(String name); + public abstract PyImportElement createImportElement(final LanguageLevel languageLevel, String name); @NotNull public abstract T createFromText(LanguageLevel langLevel, Class aClass, final String text); diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyFile.java b/python/psi-api/src/com/jetbrains/python/psi/PyFile.java index 28c3c2083aa0..05e0a940dd3a 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/PyFile.java +++ b/python/psi-api/src/com/jetbrains/python/psi/PyFile.java @@ -62,4 +62,9 @@ public interface PyFile extends PyElement, PsiFile, PyDocStringOwner, ScopeOwner * @return the deprecation message or null if the function is not deprecated. */ String getDeprecationMessage(); + + /** + * Returns the sequential list of import statements in the beginning of the file. + */ + List getImportBlock(); } diff --git a/python/src/com/jetbrains/python/PythonFoldingBuilder.java b/python/src/com/jetbrains/python/PythonFoldingBuilder.java index 2ad1d3253c53..6346c6e6a780 100644 --- a/python/src/com/jetbrains/python/PythonFoldingBuilder.java +++ b/python/src/com/jetbrains/python/PythonFoldingBuilder.java @@ -10,11 +10,12 @@ import com.intellij.openapi.util.TextRange; import com.intellij.openapi.util.text.LineTokenizer; import com.intellij.openapi.util.text.StringUtil; import com.intellij.psi.PsiElement; -import com.intellij.psi.TokenType; import com.intellij.psi.tree.IElementType; import com.jetbrains.python.psi.PyFile; import com.jetbrains.python.psi.PyFileElementType; +import com.jetbrains.python.psi.PyImportStatementBase; import com.jetbrains.python.psi.PyStringLiteralExpression; +import com.jetbrains.python.psi.impl.PyFileImpl; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -35,24 +36,12 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw private static void appendDescriptors(ASTNode node, List descriptors) { if (node.getElementType() instanceof PyFileElementType) { - ASTNode firstImport = node.getFirstChildNode(); - while(firstImport != null && !isImport(firstImport, false)) { - firstImport = firstImport.getTreeNext(); - } - if (firstImport != null) { - ASTNode lastImport = firstImport.getTreeNext(); - while(lastImport != null && isImport(lastImport.getTreeNext(), true)) { - lastImport = lastImport.getTreeNext(); - } - if (lastImport != null) { - while (lastImport.getElementType() == TokenType.WHITE_SPACE) { - lastImport = lastImport.getTreePrev(); - } - if (isImport(lastImport, false) && firstImport != lastImport) { - descriptors.add(new FoldingDescriptor(firstImport, new TextRange(firstImport.getStartOffset(), - lastImport.getTextRange().getEndOffset()))); - } - } + final List imports = ((PyFile)node.getPsi()).getImportBlock(); + if (imports.size() > 1) { + final PyImportStatementBase firstImport = imports.get(0); + final PyImportStatementBase lastImport = imports.get(imports.size()-1); + descriptors.add(new FoldingDescriptor(firstImport, new TextRange(firstImport.getTextRange().getStartOffset(), + lastImport.getTextRange().getEndOffset()))); } } else if (node.getElementType() == PyElementTypes.STATEMENT_LIST) { @@ -116,18 +105,9 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw return null; } - private static boolean isImport(ASTNode node, boolean orWhitespace) { - if (node == null) return false; - IElementType elementType = node.getElementType(); - if (orWhitespace && elementType == TokenType.WHITE_SPACE) { - return true; - } - return elementType == PyElementTypes.IMPORT_STATEMENT || elementType == PyElementTypes.FROM_IMPORT_STATEMENT; - } - @Override protected String getLanguagePlaceholderText(@NotNull ASTNode node, @NotNull TextRange range) { - if (isImport(node, false)) { + if (PyFileImpl.isImport(node, false)) { return "import ..."; } if (node.getElementType() == PyElementTypes.STRING_LITERAL_EXPRESSION) { @@ -143,7 +123,7 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw @Override protected boolean isRegionCollapsedByDefault(@NotNull ASTNode node) { - if (isImport(node, false)) { + if (PyFileImpl.isImport(node, false)) { return CodeFoldingSettings.getInstance().COLLAPSE_IMPORTS; } if (node.getElementType() == PyElementTypes.STRING_LITERAL_EXPRESSION) { diff --git a/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java b/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java index cc73160d693e..b140ee9b48a8 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java @@ -147,8 +147,9 @@ public class AddImportHelper { } } - final PyImportStatement importNodeToInsert = PyElementGenerator.getInstance(file.getProject()).createImportStatementFromText( - "import " + name + as_clause); + final PyElementGenerator generator = PyElementGenerator.getInstance(file.getProject()); + final LanguageLevel languageLevel = LanguageLevel.forElement(file); + final PyImportStatement importNodeToInsert = generator.createImportStatementFromText(languageLevel, "import " + name + as_clause); try { file.addBefore(importNodeToInsert, getInsertPosition(file, name, priority)); } @@ -167,15 +168,15 @@ public class AddImportHelper { * @param asName optional name for 'as' clause */ public static void addImportFromStatement(PsiFile file, String from, String name, @Nullable String asName, ImportPriority priority) { - String as_clause; + String asClause; if (asName == null) { - as_clause = ""; + asClause = ""; } else { - as_clause = " as " + asName; + asClause = " as " + asName; } final PyFromImportStatement importNodeToInsert = PyElementGenerator.getInstance(file.getProject()).createFromText( - LanguageLevel.getDefault(), PyFromImportStatement.class, "from " + from + " import " + name + as_clause); + LanguageLevel.forElement(file), PyFromImportStatement.class, "from " + from + " import " + name + asClause); try { file.addBefore(importNodeToInsert, getInsertPosition(file, from, priority)); } @@ -201,7 +202,8 @@ public class AddImportHelper { return false; } } - PyImportElement importElement = PyElementGenerator.getInstance(file.getProject()).createImportElement(name); + final PyElementGenerator generator = PyElementGenerator.getInstance(file.getProject()); + PyImportElement importElement = generator.createImportElement(LanguageLevel.forElement(file), name); existingImport.add(importElement); return true; } diff --git a/python/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java b/python/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java index d539ab752c69..bc70d05d84af 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java @@ -144,7 +144,7 @@ public class ImportFromExistingAction implements QuestionAction { PsiElement parent = src.getParent(); if (parent instanceof PyFromImportStatement) { // add another import element right after the one we got - PsiElement newImportElement = gen.createImportElement(myName); + PsiElement newImportElement = gen.createImportElement(LanguageLevel.getDefault(), myName); parent.add(newImportElement); } else { // just 'import' diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java index e1f1768f01ba..3ea9a670bbe4 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java @@ -2,13 +2,18 @@ package com.jetbrains.python.codeInsight.imports; import com.intellij.codeInspection.LocalInspectionToolSession; import com.intellij.lang.ImportOptimizer; +import com.intellij.psi.PsiElement; import com.intellij.psi.PsiFile; +import com.intellij.psi.PsiFileSystemItem; +import com.jetbrains.python.formatter.PyBlock; import com.jetbrains.python.inspections.PyUnresolvedReferencesInspection; -import com.jetbrains.python.psi.PyElement; -import com.jetbrains.python.psi.PyRecursiveElementVisitor; +import com.jetbrains.python.psi.*; import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; +import java.util.ArrayList; import java.util.Collections; +import java.util.List; /** * @author yole @@ -19,7 +24,7 @@ public class PyImportOptimizer implements ImportOptimizer { } @NotNull - public Runnable processFile(@NotNull PsiFile file) { + public Runnable processFile(@NotNull final PsiFile file) { final LocalInspectionToolSession session = new LocalInspectionToolSession(file, 0, file.getTextLength()); final PyUnresolvedReferencesInspection.Visitor visitor = new PyUnresolvedReferencesInspection.Visitor(null, session, @@ -34,7 +39,105 @@ public class PyImportOptimizer implements ImportOptimizer { return new Runnable() { public void run() { visitor.optimizeImports(); + if (file instanceof PyFile) { + new ImportSorter((PyFile) file).run(); + } } }; } + + private static class ImportSorter { + private final PyFile myFile; + private final List myBuiltinImports = new ArrayList(); + private final List myThirdPartyImports = new ArrayList(); + private final List myProjectImports = new ArrayList(); + private final List myImportBlock; + private final PyElementGenerator myGenerator; + private boolean myMissorted = false; + + private ImportSorter(PyFile file) { + myFile = file; + myImportBlock = myFile.getImportBlock(); + myGenerator = PyElementGenerator.getInstance(myFile.getProject()); + } + + public void run() { + if (myImportBlock.isEmpty()) { + return; + } + LanguageLevel langLevel = LanguageLevel.forElement(myFile); + for (PyImportStatementBase importStatement : myImportBlock) { + if (importStatement instanceof PyImportStatement && importStatement.getImportElements().length > 1) { + for (PyImportElement importElement : importStatement.getImportElements()) { + myMissorted = true; + PsiElement toImport = importElement.resolve(); + final PyImportStatement splitImport = myGenerator.createImportStatementFromText(langLevel, "import " + importElement.getText()); + prioritize(splitImport, toImport); + } + } + else { + PsiElement toImport; + if (importStatement instanceof PyFromImportStatement) { + toImport = ((PyFromImportStatement) importStatement).resolveImportSource(); + } + else { + toImport = importStatement.getImportElements()[0].resolve(); + } + prioritize(importStatement, toImport); + } + } + if (myMissorted) { + applyResults(); + } + } + + private void prioritize(PyImportStatementBase importStatement, @Nullable PsiElement toImport) { + if (toImport != null && !(toImport instanceof PsiFileSystemItem)) { + toImport = toImport.getContainingFile(); + } + final AddImportHelper.ImportPriority priority = toImport == null + ? AddImportHelper.ImportPriority.PROJECT + : AddImportHelper.getImportPriority(myFile, (PsiFileSystemItem)toImport); + if (priority == AddImportHelper.ImportPriority.BUILTIN) { + myBuiltinImports.add(importStatement); + if (!myThirdPartyImports.isEmpty() || !myProjectImports.isEmpty()) { + myMissorted = true; + } + } + else if (priority == AddImportHelper.ImportPriority.THIRD_PARTY) { + myThirdPartyImports.add(importStatement); + if (!myProjectImports.isEmpty()) { + myMissorted = true; + } + } + else { + myProjectImports.add(importStatement); + } + } + + private void applyResults() { + markGroupBegin(myThirdPartyImports); + markGroupBegin(myProjectImports); + addImports(myBuiltinImports); + addImports(myThirdPartyImports); + addImports(myProjectImports); + PsiElement lastElement = myImportBlock.get(myImportBlock.size()-1); + myFile.deleteChildRange(myImportBlock.get(0), lastElement); + for (PyImportStatementBase anImport : myBuiltinImports) { + anImport.putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, null); + } + } + + private static void markGroupBegin(List imports) { + if (imports.size() > 0) { + imports.get(0).putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, true); + } + } + + private void addImports(final List imports) { + for (PyImportStatementBase newImport: imports) { + myFile.addBefore(newImport, myImportBlock.get(0)); + } + } + } } diff --git a/python/src/com/jetbrains/python/formatter/PyBlock.java b/python/src/com/jetbrains/python/formatter/PyBlock.java index df745ba7706c..4f4d951213f1 100644 --- a/python/src/com/jetbrains/python/formatter/PyBlock.java +++ b/python/src/com/jetbrains/python/formatter/PyBlock.java @@ -3,6 +3,7 @@ package com.jetbrains.python.formatter; import com.intellij.formatting.*; import com.intellij.lang.ASTNode; import com.intellij.openapi.editor.Document; +import com.intellij.openapi.util.Key; import com.intellij.openapi.util.TextRange; import com.intellij.psi.*; import com.intellij.psi.impl.source.tree.TreeUtil; @@ -38,6 +39,8 @@ public class PyBlock implements ASTBlock { private Alignment myChildAlignment; private static final boolean DUMP_FORMATTING_BLOCKS = false; + public static final Key IMPORT_GROUP_BEGIN = Key.create("com.jetbrains.python.formatter.importGroupBegin"); + private static final TokenSet ourListElementTypes = TokenSet.create(PyElementTypes.LIST_LITERAL_EXPRESSION, PyElementTypes.LIST_COMP_EXPRESSION, PyElementTypes.DICT_COMP_EXPRESSION, @@ -330,6 +333,14 @@ public class PyBlock implements ASTBlock { @Nullable public Spacing getSpacing(Block child1, @NotNull Block child2) { + if (child1 instanceof ASTBlock && child2 instanceof ASTBlock) { + final PsiElement psi1 = ((ASTBlock)child1).getNode().getPsi(); + final PsiElement psi2 = ((ASTBlock)child2).getNode().getPsi(); + if (psi1 instanceof PyImportStatementBase && psi2 instanceof PyImportStatementBase && + psi2.getCopyableUserData(IMPORT_GROUP_BEGIN) != null) { + return Spacing.createSpacing(0, 0, 2, true, 1); + } + } return myContext.getSpacingBuilder().getSpacing(this, child1, child2); } diff --git a/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java b/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java index 523821335ed5..3fe0bbf99959 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java @@ -220,14 +220,15 @@ public class PyElementGeneratorImpl extends PyElementGenerator { throw new IllegalArgumentException("Invalid call expression text " + functionName); } - public PyImportStatement createImportStatementFromText(final String text) { - final PsiFile dummyFile = createDummyFile(LanguageLevel.getDefault(), text); + public PyImportStatement createImportStatementFromText(final LanguageLevel languageLevel, + final String text) { + final PsiFile dummyFile = createDummyFile(languageLevel, text); return (PyImportStatement)dummyFile.getFirstChild(); } @Override - public PyImportElement createImportElement(String name) { - return createFromText(LanguageLevel.getDefault(), PyImportElement.class, "from foo import " + name, new int[]{0, 6}); + public PyImportElement createImportElement(final LanguageLevel languageLevel, String name) { + return createFromText(languageLevel, PyImportElement.class, "from foo import " + name, new int[]{0, 6}); } static final int[] FROM_ROOT = new int[]{0}; diff --git a/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java index 372f3d150cce..4e23b2c6e93d 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java @@ -1,6 +1,7 @@ package com.jetbrains.python.psi.impl; import com.intellij.extapi.psi.PsiFileBase; +import com.intellij.lang.ASTNode; import com.intellij.lang.Language; import com.intellij.openapi.fileTypes.FileType; import com.intellij.openapi.util.Key; @@ -9,6 +10,7 @@ import com.intellij.psi.*; import com.intellij.psi.scope.PsiScopeProcessor; import com.intellij.psi.stubs.StubElement; import com.intellij.psi.templateLanguages.TemplateLanguageFileViewProvider; +import com.intellij.psi.tree.IElementType; import com.intellij.psi.util.PsiModificationTracker; import com.intellij.reference.SoftReference; import com.intellij.util.IncorrectOperationException; @@ -661,6 +663,26 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression { return extractDeprecationMessage(); } + @Override + public List getImportBlock() { + List result = new ArrayList(); + ASTNode firstImport = getNode().getFirstChildNode(); + while(firstImport != null && !isImport(firstImport, false)) { + firstImport = firstImport.getTreeNext(); + } + if (firstImport != null) { + result.add(firstImport.getPsi(PyImportStatementBase.class)); + ASTNode lastImport = firstImport.getTreeNext(); + while(lastImport != null && isImport(lastImport.getTreeNext(), true)) { + if (isImport(lastImport, false)) { + result.add(lastImport.getPsi(PyImportStatementBase.class)); + } + lastImport = lastImport.getTreeNext(); + } + } + return result; + } + public String extractDeprecationMessage() { return PyFunctionImpl.extractDeprecationMessage(getStatements()); } @@ -722,4 +744,13 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression { return new ArrayList(); } } + + public static boolean isImport(ASTNode node, boolean orWhitespace) { + if (node == null) return false; + IElementType elementType = node.getElementType(); + if (orWhitespace && elementType == TokenType.WHITE_SPACE) { + return true; + } + return elementType == PyElementTypes.IMPORT_STATEMENT || elementType == PyElementTypes.FROM_IMPORT_STATEMENT; + } } diff --git a/python/testData/folding/importBlock.py b/python/testData/folding/importBlock.py new file mode 100644 index 000000000000..a4e8699c02e6 --- /dev/null +++ b/python/testData/folding/importBlock.py @@ -0,0 +1,7 @@ +"""This is a module-level docstring.""" + +import os +import sys + +print os.path +print sys.name diff --git a/python/testData/optimizeImports/order.after.py b/python/testData/optimizeImports/order.after.py new file mode 100644 index 000000000000..e89fc2619475 --- /dev/null +++ b/python/testData/optimizeImports/order.after.py @@ -0,0 +1,9 @@ +import sys +import datetime + +import foo +from bar import * + + +sys.path +datetime.datetime diff --git a/python/testData/optimizeImports/order.py b/python/testData/optimizeImports/order.py new file mode 100644 index 000000000000..4232fbd83ed6 --- /dev/null +++ b/python/testData/optimizeImports/order.py @@ -0,0 +1,7 @@ +import foo +import sys +from bar import * +import datetime + +sys.path +datetime.datetime diff --git a/python/testData/optimizeImports/split.after.py b/python/testData/optimizeImports/split.after.py new file mode 100644 index 000000000000..3bf30c07a9fc --- /dev/null +++ b/python/testData/optimizeImports/split.after.py @@ -0,0 +1,5 @@ +import sys +import datetime + +sys.path +datetime.time diff --git a/python/testData/optimizeImports/split.py b/python/testData/optimizeImports/split.py new file mode 100644 index 000000000000..076d98a06b98 --- /dev/null +++ b/python/testData/optimizeImports/split.py @@ -0,0 +1,4 @@ +import sys, datetime + +sys.path +datetime.time diff --git a/python/testSrc/com/jetbrains/python/PyFoldingTest.java b/python/testSrc/com/jetbrains/python/PyFoldingTest.java index ec9f09f9e737..071f583b608e 100644 --- a/python/testSrc/com/jetbrains/python/PyFoldingTest.java +++ b/python/testSrc/com/jetbrains/python/PyFoldingTest.java @@ -21,4 +21,8 @@ public class PyFoldingTest extends PyTestCase { public void testCustomFolding() { doTest(); } + + public void testImportBlock() { + doTest(); + } } diff --git a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java b/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java index ab05d8be2a76..0b90b843b47c 100644 --- a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java +++ b/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java @@ -38,7 +38,15 @@ public class PyOptimizeImportsTest extends PyTestCase { public void testSuppressed() { // PY-5228 doTest(); - } + } + + public void testSplit() { + doTest(); + } + + public void testOrder() { + doTest(); + } private void doTest() { myFixture.configureByFile("optimizeImports/" + getTestName(true) + ".py");