diff --git a/python/src/com/jetbrains/python/actions/AddImportAction.java b/python/src/com/jetbrains/python/actions/AddImportAction.java index f201081a1673..a2aafc547458 100644 --- a/python/src/com/jetbrains/python/actions/AddImportAction.java +++ b/python/src/com/jetbrains/python/actions/AddImportAction.java @@ -156,8 +156,10 @@ public class AddImportAction implements HintAction, QuestionAction, LocalQuickFi @Override protected void run(Result result) throws Throwable { String name = getRefName(); - if (ResolveImportUtil.resolveInRoots(file, name) != null) { // TODO: think about multiple possible resole results - AddImportHelper.addImportStatement(file, name, null); + final PsiElement element = ResolveImportUtil.resolveInRoots(file, name); + if (element != null) { // TODO: think about multiple possible resole results + final AddImportHelper.ImportPriority priority = AddImportHelper.getImportPriority(file, element.getContainingFile()); + AddImportHelper.addImportStatement(file, name, null, priority); } } }.execute(); diff --git a/python/src/com/jetbrains/python/actions/AddImportHelper.java b/python/src/com/jetbrains/python/actions/AddImportHelper.java index 16256bf47e3e..b3e2e340fab6 100644 --- a/python/src/com/jetbrains/python/actions/AddImportHelper.java +++ b/python/src/com/jetbrains/python/actions/AddImportHelper.java @@ -1,6 +1,13 @@ package com.jetbrains.python.actions; import com.intellij.openapi.diagnostic.Logger; +import com.intellij.openapi.module.Module; +import com.intellij.openapi.module.ModuleUtil; +import com.intellij.openapi.projectRoots.Sdk; +import com.intellij.openapi.roots.OrderRootType; +import com.intellij.openapi.roots.ProjectRootManager; +import com.intellij.openapi.vfs.VfsUtil; +import com.intellij.openapi.vfs.VirtualFile; import com.intellij.psi.*; import com.intellij.util.IncorrectOperationException; import com.jetbrains.python.PythonDocStringFinder; @@ -8,6 +15,8 @@ import com.jetbrains.python.codeInsight.PyCodeInsightSettings; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyQualifiedName; import com.jetbrains.python.psi.resolve.ResolveImportUtil; +import com.jetbrains.python.psi.stubs.PyClassNameIndex; +import com.jetbrains.python.sdk.PythonSdkType; import org.jetbrains.annotations.Nullable; import java.util.List; @@ -23,7 +32,11 @@ public class AddImportHelper { private AddImportHelper() { } - private static PsiElement getInsertPosition(final PsiFile file) { + public enum ImportPriority { + BUILTIN, THIRD_PARTY, PROJECT + } + + private static PsiElement getInsertPosition(final PsiFile file, String nameToImport, ImportPriority priority) { PsiElement feeler = file.getFirstChild(); LOG.assertTrue(feeler != null); // skip initial comments and whitespace and try to get just below the last import stmt @@ -31,7 +44,10 @@ public class AddImportHelper { boolean skipped_over_doc = false; PsiElement seeker = feeler; do { - if (PyUtil.instanceOf(feeler, PyImportStatement.class, PyFromImportStatement.class)) { + if (feeler instanceof PyImportStatementBase) { + if (shouldInsertBefore(file, (PyImportStatementBase)feeler, nameToImport, priority)) { + break; + } seeker = feeler; feeler = feeler.getNextSibling(); skipped_over_imports = true; @@ -55,21 +71,49 @@ public class AddImportHelper { return seeker; } + private static boolean shouldInsertBefore(PsiFile file, PyImportStatementBase relativeTo, String nameToImport, ImportPriority priority) { + PyQualifiedName relativeToName; + PsiElement source; + if (relativeTo instanceof PyFromImportStatement) { + final PyFromImportStatement fromImportStatement = (PyFromImportStatement)relativeTo; + relativeToName = fromImportStatement.getImportSourceQName(); + source = ResolveImportUtil.resolveFromImportStatementSource(fromImportStatement); + } + else { + final PyImportElement[] importElements = relativeTo.getImportElements(); + if (importElements.length == 0) { + return false; + } + relativeToName = importElements [0].getImportedQName(); + source = ResolveImportUtil.resolveImportElement(importElements [0]); + } + if (relativeToName == null) { + return false; + } + ImportPriority relativeToPriority = source == null ? ImportPriority.BUILTIN : getImportPriority(file, source.getContainingFile()); + if (priority.compareTo(relativeToPriority) < 0) { + return true; + } + if (nameToImport.compareTo(relativeToName.toString()) < 0) { + return priority.compareTo(relativeToPriority) >= 0; + } + return false; + } + /** * Adds an import statement, presumably below all other initial imports in the file. * @param file where to operate * @param name which to import (qualified is OK) * @param asName optional name for 'as' clause - * @param project to which the file presumably belongs */ - public static void addImportStatement(PsiFile file, String name, @Nullable String asName) { + public static void addImportStatement(PsiFile file, String name, @Nullable String asName, ImportPriority priority) { String as_clause; if (asName == null) as_clause = ""; else as_clause = " as " + asName; final PyImportStatement importNodeToInsert = PyElementGenerator.getInstance(file.getProject()).createImportStatementFromText( "import " + name + as_clause); try { - file.addBefore(importNodeToInsert, getInsertPosition(file)); + file.addBefore(importNodeToInsert, getInsertPosition(file, name, priority)); } catch (IncorrectOperationException e) { LOG.error(e); @@ -78,27 +122,27 @@ public class AddImportHelper { /** * Adds an "import ... from ..." statement below other top-level imports. + * * @param file where to operate * @param from name of the module * @param name imported name * @param asName optional name for 'as' clause - * @param project where the file belongs */ - public static void addImportFromStatement(PsiFile file, String from, String name, @Nullable String asName) { + public static void addImportFromStatement(PsiFile file, String from, String name, @Nullable String asName, ImportPriority priority) { String as_clause; if (asName == null) as_clause = ""; else as_clause = " as " + asName; final PyFromImportStatement importNodeToInsert = PyElementGenerator.getInstance(file.getProject()).createFromText( LanguageLevel.getDefault(), PyFromImportStatement.class, "from " + from + " import " + name + as_clause + "\n\n"); try { - file.addBefore(importNodeToInsert, getInsertPosition(file)); + file.addBefore(importNodeToInsert, getInsertPosition(file, from, priority)); } catch (IncorrectOperationException e) { LOG.error(e); } } - public static void addImportFrom(PsiFile file, String path, final String name) { + public static void addImportFrom(PsiFile file, String path, final String name, ImportPriority priority) { final List existingImports = ((PyFile)file).getFromImports(); for (PyFromImportStatement existingImport : existingImports) { final PyQualifiedName qName = existingImport.getImportSourceQName(); @@ -108,21 +152,42 @@ public class AddImportHelper { return; } } - addImportFromStatement(file, path, name, null); + addImportFromStatement(file, path, name, null, priority); } public static void addImport(final PsiNamedElement target, final PsiFile file, final PyElement element) { final boolean useQualified = !PyCodeInsightSettings.getInstance().PREFER_FROM_IMPORT; + final ImportPriority priority = getImportPriority(file, target.getContainingFile()); final String path = ResolveImportUtil.findShortestImportableName(element, target.getContainingFile().getVirtualFile()); if (target instanceof PsiFileSystemItem) { - addImportStatement(file, path, null); + addImportStatement(file, path, null, priority); } else if (useQualified) { - addImportStatement(file, path, null); + addImportStatement(file, path, null, priority); final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(file.getProject()); element.replace(elementGenerator.createExpressionFromText(path + "." + target.getName())); } else { - addImportFrom(file, path, target.getName()); + addImportFrom(file, path, target.getName(), priority); } } + + public static ImportPriority getImportPriority(PsiElement importLocation, PsiFile fileToImport) { + final VirtualFile vFile = fileToImport.getVirtualFile(); + if (vFile == null) { + return ImportPriority.PROJECT; + } + final ProjectRootManager projectRootManager = ProjectRootManager.getInstance(fileToImport.getProject()); + if (projectRootManager.getFileIndex().isInContent(vFile)) { + return ImportPriority.PROJECT; + } + Module module = ModuleUtil.findModuleForPsiElement(importLocation); + Sdk pythonSdk = module != null ? PythonSdkType.findPythonSdk(module) : projectRootManager.getProjectSdk(); + if (pythonSdk != null) { + final VirtualFile libDir = PyClassNameIndex.findLibDir(pythonSdk.getRootProvider().getFiles(OrderRootType.CLASSES)); + if (libDir != null && VfsUtil.isAncestor(libDir, vFile, false)) { + return ImportPriority.BUILTIN; + } + } + return ImportPriority.THIRD_PARTY; + } } diff --git a/python/src/com/jetbrains/python/actions/ImportFromExistingAction.java b/python/src/com/jetbrains/python/actions/ImportFromExistingAction.java index 267b60e0b4d3..b7d700c11095 100644 --- a/python/src/com/jetbrains/python/actions/ImportFromExistingAction.java +++ b/python/src/com/jetbrains/python/actions/ImportFromExistingAction.java @@ -125,15 +125,16 @@ public class ImportFromExistingAction implements QuestionAction { } } else { // no existing import, add it then use it + AddImportHelper.ImportPriority priority = AddImportHelper.getImportPriority(myTarget, item.getFile()); if (myUseQualifiedImport) { - AddImportHelper.addImportStatement(myTarget.getContainingFile(), item.getPath(), null); + AddImportHelper.addImportStatement(myTarget.getContainingFile(), item.getPath(), null, priority); String qual_name; if (item.getAsName() != null) qual_name = item.getAsName(); else qual_name = item.getPath(); myTarget.replace(gen.createExpressionFromText(qual_name + "." + myName)); } else { - AddImportHelper.addImportFrom(myTarget.getContainingFile(), item.getPath(), myName); + AddImportHelper.addImportFrom(myTarget.getContainingFile(), item.getPath(), myName, priority); } } } diff --git a/python/src/com/jetbrains/python/psi/stubs/PyClassNameIndex.java b/python/src/com/jetbrains/python/psi/stubs/PyClassNameIndex.java index 3e0609a6db02..cbeb7959dac0 100644 --- a/python/src/com/jetbrains/python/psi/stubs/PyClassNameIndex.java +++ b/python/src/com/jetbrains/python/psi/stubs/PyClassNameIndex.java @@ -87,7 +87,7 @@ public class PyClassNameIndex extends StringStubIndexExtension { } @Nullable - private static VirtualFile findLibDir(VirtualFile[] files) { + public static VirtualFile findLibDir(VirtualFile[] files) { for (VirtualFile file : files) { if ((file.findChild("__future__.py") != null || file.findChild("__future__.pyc") != null) && file.findChild("test") != null) { return file; diff --git a/python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java b/python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java index 368299704858..213143e0e21a 100644 --- a/python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java +++ b/python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java @@ -183,15 +183,16 @@ public class PyClassRefactoringUtil { final PsiFile file = target.getContainingFile(); if (newFile == file) return; final String importableName = ResolveImportUtil.findShortestImportableName(target, vFile); + final AddImportHelper.ImportPriority priority = AddImportHelper.getImportPriority(target, newFile); if (!PyCodeInsightSettings.getInstance().PREFER_FROM_IMPORT || newClass instanceof PyFile) { if (newClass instanceof PyFile) { - AddImportHelper.addImportStatement(file, importableName, null); + AddImportHelper.addImportStatement(file, importableName, null, priority); } else { final String name = newClass.getName(); - AddImportHelper.addImportStatement(file, importableName + "." + name, null); + AddImportHelper.addImportStatement(file, importableName + "." + name, null, priority); } } else { - AddImportHelper.addImportFrom(file, importableName, newClass.getName()); + AddImportHelper.addImportFrom(file, importableName, newClass.getName(), priority); } } diff --git a/python/testData/addImport/addBuiltin.after.py b/python/testData/addImport/addBuiltin.after.py new file mode 100644 index 000000000000..d598f9ec31ad --- /dev/null +++ b/python/testData/addImport/addBuiltin.after.py @@ -0,0 +1,3 @@ +import datetime +import re +import sys diff --git a/python/testData/addImport/addBuiltin.py b/python/testData/addImport/addBuiltin.py new file mode 100644 index 000000000000..71f385992d2b --- /dev/null +++ b/python/testData/addImport/addBuiltin.py @@ -0,0 +1,2 @@ +import datetime +import sys diff --git a/python/testSrc/com/jetbrains/python/PyAddImportTest.java b/python/testSrc/com/jetbrains/python/PyAddImportTest.java new file mode 100644 index 000000000000..d4c96d83a18f --- /dev/null +++ b/python/testSrc/com/jetbrains/python/PyAddImportTest.java @@ -0,0 +1,22 @@ +package com.jetbrains.python; + +import com.intellij.openapi.application.Result; +import com.intellij.openapi.command.WriteCommandAction; +import com.jetbrains.python.actions.AddImportHelper; +import com.jetbrains.python.fixtures.PyLightFixtureTestCase; + +/** + * @author yole + */ +public class PyAddImportTest extends PyLightFixtureTestCase { + public void testAddBuiltin() { + myFixture.configureByFile("addImport/addBuiltin.py"); + new WriteCommandAction(myFixture.getProject(), myFixture.getFile()) { + @Override + protected void run(Result result) throws Throwable { + AddImportHelper.addImportStatement(myFixture.getFile(), "re", null, AddImportHelper.ImportPriority.BUILTIN); + } + }.execute(); + myFixture.checkResultByFile("addImport/addBuiltin.after.py"); + } +}