PY-6020 Implement feature of auto import inside functions

This commit is contained in:
Mikhail Golubev
2014-08-07 20:09:35 +04:00
parent 9e44bfaac0
commit 7576d369b4
12 changed files with 213 additions and 40 deletions
@@ -90,8 +90,6 @@ public abstract class PyElementGenerator {
@NotNull
public abstract PyCallExpression createCallExpression(final LanguageLevel langLevel, String functionName);
public abstract PyImportStatement createImportStatementFromText(final LanguageLevel languageLevel, String text);
public abstract PyImportElement createImportElement(final LanguageLevel languageLevel, String name);
public abstract PyFunction createProperty(final LanguageLevel languageLevel,
@@ -139,4 +137,32 @@ public abstract class PyElementGenerator {
*/
@NotNull
public abstract PsiElement createNewLine();
/**
* Creates import statement of form {@code from qualifier import name as alias}.
*
* @param languageLevel language level for created element
* @param qualifier from where {@code name} will be imported (module name)
* @param name text of the reference in import element
* @param alias optional alias for {@code as alias} part
* @return created {@link com.jetbrains.python.psi.PyFromImportStatement}
*/
@NotNull
public abstract PyFromImportStatement createFromImportStatement(@NotNull LanguageLevel languageLevel,
@NotNull String qualifier,
@NotNull String name,
@Nullable String alias);
/**
* Creates import statement of form {@code import name as alias}.
*
* @param languageLevel language level for created element
* @param name text of the reference in import element (module name)
* @param alias optional alias for {@code as alias} part
* @return created {@link com.jetbrains.python.psi.PyImportStatement}
*/
@NotNull
public abstract PyImportStatement createImportStatement(@NotNull LanguageLevel languageLevel,
@NotNull String name,
@Nullable String alias);
}
@@ -21,14 +21,14 @@ import com.intellij.openapi.module.Module;
import com.intellij.openapi.module.ModuleUtilCore;
import com.intellij.openapi.projectRoots.Sdk;
import com.intellij.openapi.roots.ProjectRootManager;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.*;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.psi.util.QualifiedName;
import com.intellij.util.IncorrectOperationException;
import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
import com.jetbrains.python.documentation.DocStringUtil;
import com.jetbrains.python.psi.*;
import com.intellij.psi.util.QualifiedName;
import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
import com.jetbrains.python.sdk.PythonSdkType;
import org.jetbrains.annotations.NotNull;
@@ -36,6 +36,8 @@ import org.jetbrains.annotations.Nullable;
import java.util.List;
import static com.jetbrains.python.psi.PyUtil.sure;
/**
* Does the actual job of adding an import statement into a file.
* User: dcheryasov
@@ -47,6 +49,35 @@ public class AddImportHelper {
private AddImportHelper() {
}
public static void addLocalImportStatement(@NotNull PyElement element, @NotNull String name) {
final PyElementGenerator generator = PyElementGenerator.getInstance(element.getProject());
final LanguageLevel languageLevel = LanguageLevel.forElement(element);
final PsiElement anchor = getLocalInsertPosition(element);
final PsiElement parentElement = sure(anchor).getParent();
if (parentElement != null) {
parentElement.addBefore(generator.createImportStatement(languageLevel, name, null), anchor);
}
}
public static void addLocalFromImportStatement(@NotNull PyElement element, @NotNull String qualifier, @NotNull String name) {
final PyElementGenerator generator = PyElementGenerator.getInstance(element.getProject());
final LanguageLevel languageLevel = LanguageLevel.forElement(element);
final PsiElement anchor = getLocalInsertPosition(element);
final PsiElement parentElement = sure(anchor).getParent();
if (parentElement != null) {
parentElement.addBefore(generator.createFromImportStatement(languageLevel, qualifier, name, null), anchor);
}
}
@Nullable
public static PsiElement getLocalInsertPosition(@NotNull PyElement anchor) {
final PyStatement enclosingStatement = PsiTreeUtil.getParentOfType(anchor, PyStatement.class, false);
return enclosingStatement != null ? enclosingStatement.getFirstChild() : null;
}
public enum ImportPriority {
BUILTIN, THIRD_PARTY, PROJECT
}
@@ -81,7 +112,8 @@ public class AddImportHelper {
// maybe we arrived at the doc comment stmt; skip over it, too
else if (!skippedOverImports && !skippedOverDoc && file instanceof PyFile) {
PsiElement doc_elt =
DocStringUtil.findDocStringExpression((PyElement)file); // this gives the literal; its parent is the expr seeker may have encountered
DocStringUtil
.findDocStringExpression((PyElement)file); // this gives the literal; its parent is the expr seeker may have encountered
if (doc_elt != null && doc_elt.getParent() == feeler) {
feeler = feeler.getNextSibling();
seeker = feeler; // skip over doc even if there's nothing below it
@@ -147,19 +179,13 @@ public class AddImportHelper {
* @param file where to operate
* @param name which to import (qualified is OK)
* @param asName optional name for 'as' clause
* @return whether import statement was actually added
*/
public static boolean addImportStatement(PsiFile file, String name, @Nullable String asName, ImportPriority priority) {
String as_clause;
if (asName == null) {
as_clause = "";
}
else {
as_clause = " as " + asName;
}
if (!(file instanceof PyFile)) {
return false;
}
List<PyImportElement> existingImports = ((PyFile)file).getImportTargets();
final List<PyImportElement> existingImports = ((PyFile)file).getImportTargets();
for (PyImportElement element : existingImports) {
final QualifiedName qName = element.getImportedQName();
if (qName != null && name.equals(qName.toString())) {
@@ -171,7 +197,7 @@ public class AddImportHelper {
final PyElementGenerator generator = PyElementGenerator.getInstance(file.getProject());
final LanguageLevel languageLevel = LanguageLevel.forElement(file);
final PyImportStatement importNodeToInsert = generator.createImportStatementFromText(languageLevel, "import " + name + as_clause);
final PyImportStatement importNodeToInsert = generator.createImportStatement(languageLevel, name, asName);
try {
file.addBefore(importNodeToInsert, getInsertPosition(file, name, priority));
}
@@ -180,6 +206,7 @@ public class AddImportHelper {
}
return true;
}
/**
* Adds an "import ... from ..." statement below other top-level imports.
*
@@ -189,20 +216,20 @@ 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 asClause = asName == null ? "" : " as " + asName;
final PyFromImportStatement importNodeToInsert = PyElementGenerator.getInstance(file.getProject()).createFromText(
LanguageLevel.forElement(file), PyFromImportStatement.class, "from " + from + " import " + name + asClause);
final PyElementGenerator generator = PyElementGenerator.getInstance(file.getProject());
final LanguageLevel languageLevel = LanguageLevel.forElement(file);
final PyFromImportStatement nodeToInsert = generator.createFromImportStatement(languageLevel, from, name, asName);
try {
if (InjectedLanguageManager.getInstance(file.getProject()).isInjectedFragment(file)) {
final PsiElement element = file.addBefore(importNodeToInsert, getInsertPosition(file, from, priority));
final PsiElement element = file.addBefore(nodeToInsert, getInsertPosition(file, from, priority));
PsiElement whitespace = element.getNextSibling();
if (!(whitespace instanceof PsiWhiteSpace))
if (!(whitespace instanceof PsiWhiteSpace)) {
whitespace = PsiParserFacade.SERVICE.getInstance(file.getProject()).createWhiteSpaceFromText(" >>> ");
}
file.addBefore(whitespace, element);
}
else {
file.addBefore(importNodeToInsert, getInsertPosition(file, from, priority));
file.addBefore(nodeToInsert, getInsertPosition(file, from, priority));
}
}
catch (IncorrectOperationException e) {
@@ -228,7 +255,7 @@ public class AddImportHelper {
}
}
final PyElementGenerator generator = PyElementGenerator.getInstance(file.getProject());
PyImportElement importElement = generator.createImportElement(LanguageLevel.forElement(file), name);
final PyImportElement importElement = generator.createImportElement(LanguageLevel.forElement(file), name);
existingImport.add(importElement);
return true;
}
@@ -239,7 +266,8 @@ public class AddImportHelper {
public static void addImport(final PsiNamedElement target, final PsiFile file, final PyElement element) {
final boolean useQualified = !PyCodeInsightSettings.getInstance().PREFER_FROM_IMPORT;
final PsiFileSystemItem toImport = target instanceof PsiFileSystemItem ? ((PsiFileSystemItem)target).getParent() : target.getContainingFile();
final PsiFileSystemItem toImport =
target instanceof PsiFileSystemItem ? ((PsiFileSystemItem)target).getParent() : target.getContainingFile();
final ImportPriority priority = getImportPriority(file, toImport);
final QualifiedName qName = QualifiedNameFinder.findCanonicalImportPath(target, element);
if (qName == null) return;
@@ -29,6 +29,7 @@ import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiFileSystemItem;
import com.intellij.psi.PsiReference;
import com.intellij.psi.util.QualifiedName;
import com.intellij.util.IncorrectOperationException;
import com.jetbrains.python.PyBundle;
import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
@@ -36,7 +37,6 @@ import com.jetbrains.python.psi.PyElement;
import com.jetbrains.python.psi.PyFunction;
import com.jetbrains.python.psi.PyImportElement;
import com.jetbrains.python.psi.PyQualifiedExpression;
import com.intellij.psi.util.QualifiedName;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
@@ -138,7 +138,7 @@ public class AutoImportQuickFix implements LocalQuickFix, HighPriorityAction {
myImports.size() > 1,
ImportCandidateHolder.getQualifiedName(name, myImports.get(0).getPath(), myImports.get(0).getImportElement())
);
final ImportFromExistingAction action = new ImportFromExistingAction(myNode, myImports, name, myUseQualifiedImport);
final ImportFromExistingAction action = new ImportFromExistingAction(myNode, myImports, name, myUseQualifiedImport, false);
action.onDone(new Runnable() {
public void run() {
myExpended = true;
@@ -166,11 +166,16 @@ public class AutoImportQuickFix implements LocalQuickFix, HighPriorityAction {
if (!FileModificationService.getInstance().prepareFileForWrite(file)) return;
if (ImportFromExistingAction.isResolved(myReference)) return;
// act
ImportFromExistingAction action = new ImportFromExistingAction(myNode, myImports, getNameToImport(), myUseQualifiedImport);
ImportFromExistingAction action = createAction();
action.execute(); // assume that action runs in WriteAction on its own behalf
myExpended = true;
}
@NotNull
protected ImportFromExistingAction createAction() {
return new ImportFromExistingAction(myNode, myImports, getNameToImport(), myUseQualifiedImport, false);
}
public void sortCandidates() {
Collections.sort(myImports);
}
@@ -203,4 +208,27 @@ public class AutoImportQuickFix implements LocalQuickFix, HighPriorityAction {
}
return false;
}
@NotNull
public AutoImportQuickFix forLocalImport() {
return new AutoImportQuickFix(myNode, myReference, myUseQualifiedImport) {
@NotNull
@Override
public String getName() {
return super.getName() + " locally";
}
@NotNull
@Override
public String getFamilyName() {
return "import locally";
}
@NotNull
@Override
protected ImportFromExistingAction createAction() {
return new ImportFromExistingAction(myNode, myImports, getNameToImport(), myUseQualifiedImport, true);
}
};
}
}
@@ -53,6 +53,7 @@ public class ImportFromExistingAction implements QuestionAction {
String myName;
boolean myUseQualifiedImport;
private Runnable myOnDoneCallback;
private final boolean myImportLocally;
/**
* @param target element to become qualified as imported.
@@ -60,12 +61,13 @@ public class ImportFromExistingAction implements QuestionAction {
* @param name relevant name ot the target element (e.g. of identifier in an expression).
* @param useQualified if True, use qualified "import modulename" instead of "from modulename import ...".
*/
public ImportFromExistingAction(@NotNull PyElement target, @NotNull List<ImportCandidateHolder> sources, String name,
boolean useQualified) {
public ImportFromExistingAction(@NotNull PyElement target, @NotNull List<ImportCandidateHolder> sources, @NotNull String name,
boolean useQualified, boolean importLocally) {
myTarget = target;
mySources = sources;
myName = name;
myUseQualifiedImport = useQualified;
myImportLocally = importLocally;
}
public void onDone(Runnable callback) {
@@ -151,25 +153,41 @@ public class ImportFromExistingAction implements QuestionAction {
if (manager.isInjectedFragment(file)) {
file = manager.getTopLevelFile(myTarget);
}
// We are trying to import top-level module or package which thus cannot be qualified
if (isRoot(item.getFile())) {
AddImportHelper.addImportStatement(file, myName, null, priority);
if (myImportLocally) {
AddImportHelper.addLocalImportStatement(myTarget, myName);
} else {
AddImportHelper.addImportStatement(file, myName, null, priority);
}
}
else {
String qualifiedName = item.getPath().toString();
final String qualifiedName = item.getPath().toString();
if (myUseQualifiedImport) {
String nameToImport = qualifiedName;
if (item.getImportable() instanceof PsiFileSystemItem) {
nameToImport += "." + myName;
}
AddImportHelper.addImportStatement(file, nameToImport, null, priority);
if (myImportLocally) {
AddImportHelper.addLocalImportStatement(myTarget, nameToImport);
}
else {
AddImportHelper.addImportStatement(file, nameToImport, null, priority);
}
myTarget.replace(gen.createExpressionFromText(LanguageLevel.forElement(myTarget), qualifiedName + "." + myName));
}
else {
AddImportHelper.addImportFrom(file, myTarget, qualifiedName, myName, null, priority);
if (myImportLocally) {
AddImportHelper.addLocalFromImportStatement(myTarget, qualifiedName, myName);
}
else {
AddImportHelper.addImportFromStatement(file, qualifiedName, myName, null, priority);
}
}
}
}
private void addToExistingImport(PyImportElement src) {
final PyElementGenerator gen = PyElementGenerator.getInstance(myTarget.getProject());
// did user choose 'import' or 'from import'?
@@ -92,7 +92,7 @@ public class PyImportOptimizer implements ImportOptimizer {
for (PyImportElement importElement : importStatement.getImportElements()) {
myMissorted = true;
PsiElement toImport = importElement.resolve();
final PyImportStatement splitImport = myGenerator.createImportStatementFromText(langLevel, "import " + importElement.getText());
final PyImportStatement splitImport = myGenerator.createImportStatement(langLevel, importElement.getText(), null);
prioritize(splitImport, toImport);
}
}
@@ -69,7 +69,6 @@ import org.jetbrains.annotations.Nullable;
import javax.swing.*;
import java.util.*;
import java.util.HashSet;
import static com.jetbrains.python.inspections.quickfix.AddIgnoredIdentifierQuickFix.END_WILDCARD;
@@ -833,6 +832,9 @@ public class PyUnresolvedReferencesInspection extends PyInspection {
else {
actions.add(importFix);
}
if (ScopeUtil.getScopeOwner(node) instanceof PyFunction) {
actions.add(importFix.forLocalImport());
}
}
}
@@ -20,6 +20,7 @@ import com.google.common.collect.Queues;
import com.intellij.lang.ASTNode;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.Pair;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
@@ -285,12 +286,6 @@ public class PyElementGeneratorImpl extends PyElementGenerator {
throw new IllegalArgumentException("Invalid call expression text " + functionName);
}
public PyImportStatement createImportStatementFromText(final LanguageLevel languageLevel,
final String text) {
final PsiFile dummyFile = createDummyFile(languageLevel, text);
return (PyImportStatement)dummyFile.getFirstChild();
}
@Override
public PyImportElement createImportElement(final LanguageLevel languageLevel, String name) {
return createFromText(languageLevel, PyImportElement.class, "from foo import " + name, new int[]{0, 6});
@@ -427,6 +422,23 @@ public class PyElementGeneratorImpl extends PyElementGenerator {
return createFromText(LanguageLevel.getDefault(), PsiWhiteSpace.class, " \n\n ");
}
@NotNull
@Override
public PyFromImportStatement createFromImportStatement(@NotNull LanguageLevel languageLevel, @NotNull String qualifier,
@NotNull String name, @Nullable String alias) {
final String asClause = StringUtil.isNotEmpty(alias) ? " as " + alias : "";
final String statement = "from " + qualifier + " import " + name + asClause;
return createFromText(languageLevel, PyFromImportStatement.class, statement);
}
@NotNull
@Override
public PyImportStatement createImportStatement(@NotNull LanguageLevel languageLevel, @NotNull String name, @Nullable String alias) {
final String asClause = StringUtil.isNotEmpty(alias) ? " as " + alias : "";
final String statement = "import " + name + asClause;
return createFromText(languageLevel, PyImportStatement.class, statement);
}
private static class CommasOnly extends NotNullPredicate<LeafPsiElement> {
@Override
protected boolean applyNotNull(@NotNull final LeafPsiElement input) {
@@ -0,0 +1,5 @@
def func():
for _ range(10):
from package.module import foo
foo
# <ref>
@@ -0,0 +1,4 @@
def func():
for _ range(10):
foo
# <ref>
@@ -0,0 +1,7 @@
def func():
try:
import module
module
# <ref>
except:
pass
+6
View File
@@ -0,0 +1,6 @@
def func():
try:
module
# <ref>
except:
pass
@@ -17,8 +17,13 @@ package com.jetbrains.python;
import com.intellij.openapi.application.Result;
import com.intellij.openapi.command.WriteCommandAction;
import com.intellij.psi.PsiPolyVariantReference;
import com.jetbrains.python.codeInsight.imports.AddImportHelper;
import com.jetbrains.python.fixtures.PyResolveTestCase;
import com.jetbrains.python.fixtures.PyTestCase;
import com.jetbrains.python.psi.PyElement;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
/**
* @author yole
@@ -53,4 +58,36 @@ public class PyAddImportTest extends PyTestCase {
}.execute();
myFixture.checkResultByFile("addImport/" + getTestName(true) + ".after.py");
}
// PY-6020
public void testLocalFromImport() {
doAddLocalImport("foo", "package.module");
}
// PY-6020
public void testLocalImport() {
doAddLocalImport("module", null);
}
/**
* Add local import statement
* @param name reference name in corresponding import element
* @param qualifier if not {@code null} form {@code from qualifier import name} will be used, otherwise {@code import name}
*/
private void doAddLocalImport(@NotNull final String name, @Nullable final String qualifier) {
myFixture.configureByFile("addImport/" + getTestName(true) + ".py");
new WriteCommandAction(myFixture.getProject(), myFixture.getFile()) {
@Override
protected void run(Result result) throws Throwable {
final PsiPolyVariantReference reference = PyResolveTestCase.findReferenceByMarker(myFixture.getFile());
if (qualifier != null) {
AddImportHelper.addLocalFromImportStatement((PyElement)reference.getElement(), qualifier, name);
}
else {
AddImportHelper.addLocalImportStatement((PyElement)reference.getElement(), name);
}
}
}.execute();
myFixture.checkResultByFile("addImport/" + getTestName(true) + ".after.py");
}
}