From 9df1dacb23a4a7126e9aa1d6ba8302fc72361321 Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Fri, 27 May 2016 19:03:44 +0300 Subject: [PATCH] PyFile#getImportBlock works better when there are comments in-between or file ends with import --- .../python/PythonFoldingBuilder.java | 10 ++++-- .../jetbrains/python/psi/impl/PyFileImpl.java | 36 ++++++++----------- .../extractImportBlockNoWhitespaceAtEnd.py | 2 ++ ...ractImportBlockWithIntermediateComments.py | 5 +++ .../python/PyOptimizeImportsTest.java | 18 ++++++++++ 5 files changed, 46 insertions(+), 25 deletions(-) create mode 100644 python/testData/optimizeImports/extractImportBlockNoWhitespaceAtEnd.py create mode 100644 python/testData/optimizeImports/extractImportBlockWithIntermediateComments.py diff --git a/python/src/com/jetbrains/python/PythonFoldingBuilder.java b/python/src/com/jetbrains/python/PythonFoldingBuilder.java index 71f22aa670e2..0b9c8b375233 100644 --- a/python/src/com/jetbrains/python/PythonFoldingBuilder.java +++ b/python/src/com/jetbrains/python/PythonFoldingBuilder.java @@ -30,7 +30,6 @@ import com.intellij.psi.PsiWhiteSpace; import com.intellij.psi.tree.IElementType; import com.intellij.psi.tree.TokenSet; import com.jetbrains.python.psi.*; -import com.jetbrains.python.psi.impl.PyFileImpl; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -201,7 +200,7 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw @Override protected String getLanguagePlaceholderText(@NotNull ASTNode node, @NotNull TextRange range) { - if (PyFileImpl.isImport(node, false)) { + if (isImport(node)) { return "import ..."; } if (node.getElementType() == PyElementTypes.STRING_LITERAL_EXPRESSION) { @@ -231,7 +230,7 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw @Override protected boolean isRegionCollapsedByDefault(@NotNull ASTNode node) { - if (PyFileImpl.isImport(node, false)) { + if (isImport(node)) { return CodeFoldingSettings.getInstance().COLLAPSE_IMPORTS; } if (node.getElementType() == PyElementTypes.STRING_LITERAL_EXPRESSION) { @@ -265,4 +264,9 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw protected boolean isCustomFoldingRoot(ASTNode node) { return node.getPsi() instanceof PyFile || node.getElementType() == PyElementTypes.STATEMENT_LIST; } + + private static boolean isImport(@NotNull ASTNode node) { + return PyElementTypes.IMPORT_STATEMENTS.contains(node.getElementType()); + } + } diff --git a/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java index 26fc13e1c796..547812910282 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java @@ -19,7 +19,6 @@ import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.intellij.extapi.psi.PsiFileBase; import com.intellij.icons.AllIcons; -import com.intellij.lang.ASTNode; import com.intellij.lang.Language; import com.intellij.navigation.ItemPresentation; import com.intellij.openapi.fileTypes.FileType; @@ -33,8 +32,8 @@ import com.intellij.openapi.vfs.VirtualFile; import com.intellij.psi.*; import com.intellij.psi.scope.PsiScopeProcessor; import com.intellij.psi.stubs.StubElement; -import com.intellij.psi.tree.IElementType; import com.intellij.psi.util.PsiModificationTracker; +import com.intellij.psi.util.PsiTreeUtil; import com.intellij.psi.util.QualifiedName; import com.intellij.reference.SoftReference; import com.intellij.util.IncorrectOperationException; @@ -624,19 +623,21 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression { @Override public List getImportBlock() { - List result = new ArrayList(); - ASTNode firstImport = getNode().getFirstChildNode(); - while (firstImport != null && !isImport(firstImport, false)) { - firstImport = firstImport.getTreeNext(); + final List result = new ArrayList<>(); + final PsiElement firstChild = getFirstChild(); + final PyImportStatementBase firstImport; + if (firstChild instanceof PyImportStatementBase) { + firstImport = (PyImportStatementBase)firstChild; + } + else { + firstImport = PsiTreeUtil.getNextSiblingOfType(firstChild, PyImportStatementBase.class); } 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(); + result.add(firstImport); + PsiElement nextImport = PyPsiUtils.getNextNonCommentSibling(firstImport, true); + while (nextImport instanceof PyImportStatementBase) { + result.add((PyImportStatementBase)nextImport); + nextImport = PyPsiUtils.getNextNonCommentSibling(nextImport, true); } } return result; @@ -729,15 +730,6 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression { } } - 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; - } - @Override public ItemPresentation getPresentation() { return new ItemPresentation() { diff --git a/python/testData/optimizeImports/extractImportBlockNoWhitespaceAtEnd.py b/python/testData/optimizeImports/extractImportBlockNoWhitespaceAtEnd.py new file mode 100644 index 000000000000..eadf8fffe7e6 --- /dev/null +++ b/python/testData/optimizeImports/extractImportBlockNoWhitespaceAtEnd.py @@ -0,0 +1,2 @@ +import foo +import bar \ No newline at end of file diff --git a/python/testData/optimizeImports/extractImportBlockWithIntermediateComments.py b/python/testData/optimizeImports/extractImportBlockWithIntermediateComments.py new file mode 100644 index 000000000000..07439c8ca7ef --- /dev/null +++ b/python/testData/optimizeImports/extractImportBlockWithIntermediateComments.py @@ -0,0 +1,5 @@ +import sys +# noinspection PyUnresolvedReferences +import typing + +print(sys.path) diff --git a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java b/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java index 3226eb60ae5a..012df3c0350b 100644 --- a/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java +++ b/python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java @@ -23,8 +23,12 @@ import com.intellij.openapi.projectRoots.SdkModificator; import com.intellij.openapi.roots.OrderRootType; import com.intellij.openapi.vfs.VirtualFile; import com.jetbrains.python.fixtures.PyTestCase; +import com.jetbrains.python.psi.PyImportStatementBase; +import com.jetbrains.python.psi.impl.PyFileImpl; import com.jetbrains.python.sdk.PythonSdkType; +import java.util.List; + /** * @author yole */ @@ -191,6 +195,20 @@ public class PyOptimizeImportsTest extends PyTestCase { myFixture.checkResultByFile(testName + "/pkg/main.after.py"); } + public void testExtractImportBlockWithIntermediateComments() { + myFixture.configureByFile(getTestName(true) + ".py"); + final PyFileImpl file = assertInstanceOf(myFixture.getFile(), PyFileImpl.class); + final List block = file.getImportBlock(); + assertSize(2, block); + } + + public void testExtractImportBlockNoWhitespaceAtEnd() { + myFixture.configureByFile(getTestName(true) + ".py"); + final PyFileImpl file = assertInstanceOf(myFixture.getFile(), PyFileImpl.class); + final List block = file.getImportBlock(); + assertSize(2, block); + } + private void doTest() { myFixture.configureByFile(getTestName(true) + ".py"); OptimizeImportsAction.actionPerformedImpl(DataManager.getInstance().getDataContext(myFixture.getEditor().getContentComponent()));