PyFile#getImportBlock works better when there are comments in-between or file ends with import

This commit is contained in:
Mikhail Golubev
2016-06-15 19:34:13 +03:00
parent 5e045e5c12
commit 9df1dacb23
5 changed files with 46 additions and 25 deletions
@@ -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());
}
}
@@ -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<PyImportStatementBase> getImportBlock() {
List<PyImportStatementBase> result = new ArrayList<PyImportStatementBase>();
ASTNode firstImport = getNode().getFirstChildNode();
while (firstImport != null && !isImport(firstImport, false)) {
firstImport = firstImport.getTreeNext();
final List<PyImportStatementBase> 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() {
@@ -0,0 +1,2 @@
import foo
import bar
@@ -0,0 +1,5 @@
import sys
# noinspection PyUnresolvedReferences
import typing
print(sys.path)
@@ -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<PyImportStatementBase> block = file.getImportBlock();
assertSize(2, block);
}
public void testExtractImportBlockNoWhitespaceAtEnd() {
myFixture.configureByFile(getTestName(true) + ".py");
final PyFileImpl file = assertInstanceOf(myFixture.getFile(), PyFileImpl.class);
final List<PyImportStatementBase> block = file.getImportBlock();
assertSize(2, block);
}
private void doTest() {
myFixture.configureByFile(getTestName(true) + ".py");
OptimizeImportsAction.actionPerformedImpl(DataManager.getInstance().getDataContext(myFixture.getEditor().getContentComponent()));