diff --git a/platform/core-api/src/com/intellij/psi/util/QualifiedName.java b/platform/core-api/src/com/intellij/psi/util/QualifiedName.java index 7ff4f7c4fd2a..de45ed6b6f4c 100644 --- a/platform/core-api/src/com/intellij/psi/util/QualifiedName.java +++ b/platform/core-api/src/com/intellij/psi/util/QualifiedName.java @@ -31,7 +31,7 @@ import java.util.List; /** * @author yole */ -public class QualifiedName { +public class QualifiedName implements Comparable { @NotNull private final List myComponents; private QualifiedName(int count) { @@ -205,4 +205,9 @@ public class QualifiedName { public QualifiedName subQualifiedName(int fromIndex, int toIndex) { return fromComponents(myComponents.subList(fromIndex, toIndex)); } + + @Override + public int compareTo(@NotNull QualifiedName other) { + return toString().compareTo(other.toString()); + } } diff --git a/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java b/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java index e945f317f86c..df12aa70526b 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java @@ -21,14 +21,17 @@ 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.Comparing; 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.ArrayUtil; import com.intellij.util.IncorrectOperationException; import com.jetbrains.python.codeInsight.PyCodeInsightSettings; import com.jetbrains.python.documentation.DocStringUtil; +import com.jetbrains.python.formatter.PyBlock; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.resolve.QualifiedNameFinder; import com.jetbrains.python.sdk.PythonSdkType; @@ -79,7 +82,10 @@ public class AddImportHelper { } public enum ImportPriority { - BUILTIN, THIRD_PARTY, PROJECT + FUTURE, + BUILTIN, + THIRD_PARTY, + PROJECT } @Nullable @@ -88,7 +94,8 @@ public class AddImportHelper { } @Nullable - private static PsiElement getInsertPosition(final PsiElement insertParent, @Nullable String nameToImport, + private static PsiElement getInsertPosition(@NotNull PsiElement insertParent, + @Nullable PyImportStatementBase newImport, @Nullable ImportPriority priority) { PsiElement feeler = insertParent.getFirstChild(); if (feeler == null) return null; @@ -97,10 +104,18 @@ public class AddImportHelper { boolean skippedOverDoc = false; PsiElement seeker = feeler; final boolean isInjected = InjectedLanguageManager.getInstance(feeler.getProject()).isInjectedFragment(feeler.getContainingFile()); + PyImportStatementBase importAbove = null, importBelow = null; do { if (feeler instanceof PyImportStatementBase && !isInjected) { - if (nameToImport != null && priority != null && shouldInsertBefore(insertParent, (PyImportStatementBase)feeler, nameToImport, priority)) { - break; + final PyImportStatementBase existingImport = (PyImportStatementBase)feeler; + if (priority != null && newImport != null) { + if (shouldInsertBefore(newImport, existingImport, priority)) { + importBelow = existingImport; + break; + } + else { + importAbove = existingImport; + } } seeker = feeler; feeler = feeler.getNextSibling(); @@ -112,9 +127,9 @@ public class AddImportHelper { } // maybe we arrived at the doc comment stmt; skip over it, too else if (!skippedOverImports && !skippedOverDoc && insertParent instanceof PyFile) { - PsiElement doc_elt = - DocStringUtil.findDocStringExpression((PyElement)insertParent); // this gives the literal; its parent is the expr seeker may have encountered - if (doc_elt != null && doc_elt.getParent() == feeler) { + // this gives the literal; its parent is the expr seeker may have encountered + final PsiElement docElem = DocStringUtil.findDocStringExpression((PyElement)insertParent); + if (docElem != null && docElem.getParent() == feeler) { feeler = feeler.getNextSibling(); seeker = feeler; // skip over doc even if there's nothing below it skippedOverDoc = true; @@ -128,50 +143,95 @@ public class AddImportHelper { } } while (feeler != null); + final ImportPriority priorityAbove = importAbove != null ? getImportPriority(importAbove) : null; + final ImportPriority priorityBelow = importBelow != null ? getImportPriority(importBelow) : null; + if (priorityAbove != null && priority.compareTo(priorityAbove) > 0) { + newImport.putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, true); + if (priorityBelow == priority) { + importBelow.putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, null); + } + } + if (priorityBelow != null && priority.compareTo(priorityBelow) < 0) { + importBelow.putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, true); + } return seeker; } - private static boolean shouldInsertBefore(PsiElement element, PyImportStatementBase relativeTo, String nameToImport, - ImportPriority priority) { - QualifiedName relativeToName; - PsiElement source; - if (relativeTo instanceof PyFromImportStatement) { - final PyFromImportStatement fromImportStatement = (PyFromImportStatement)relativeTo; - if (fromImportStatement.isFromFuture()) { - return false; - } - relativeToName = fromImportStatement.getImportSourceQName(); - source = fromImportStatement.resolveImportSource(); + private static boolean shouldInsertBefore(@Nullable PyImportStatementBase newImport, + @NotNull PyImportStatementBase existingImport, + @NotNull ImportPriority priority) { + final ImportPriority existingImportPriority = getImportPriority(existingImport); + final int byPriority = priority.compareTo(existingImportPriority); + if (byPriority != 0) { + return byPriority < 0; } - else { - final PyImportElement[] importElements = relativeTo.getImportElements(); - if (importElements.length == 0) { - return false; - } - relativeToName = importElements[0].getImportedQName(); - source = importElements[0].resolve(); - } - if (relativeToName == null) { + if (newImport == null) { return false; } - final PsiFileSystemItem containingFile; - if (source instanceof PsiDirectory) { - containingFile = (PsiDirectory)source; + final QualifiedName existingImportName = getImportFirstQualifiedName(existingImport); + final QualifiedName newImportName = getImportFirstQualifiedName(newImport); + return Comparing.compare(newImportName, existingImportName) < 0; + } + + @Nullable + public static QualifiedName getImportFirstQualifiedName(@NotNull PyImportStatementBase importStatement) { + if (importStatement instanceof PyFromImportStatement) { + return ((PyFromImportStatement)importStatement).getImportSourceQName(); + } + else if (importStatement instanceof PyImportStatement) { + final PyImportElement importElement = ArrayUtil.getFirstElement(importStatement.getImportElements()); + if (importElement != null) { + return importElement.getImportedQName(); + } + } + return null; + } + + @NotNull + public static ImportPriority getImportPriority(@NotNull PyImportStatementBase importStatement) { + final PsiElement resolved; + if (importStatement instanceof PyFromImportStatement) { + final PyFromImportStatement fromImportStatement = (PyFromImportStatement)importStatement; + if (fromImportStatement.isFromFuture()) { + return ImportPriority.FUTURE; + } + resolved = fromImportStatement.resolveImportSource(); } else { - containingFile = source != null ? source.getContainingFile() : null; + final PyImportElement firstImportElement = ArrayUtil.getFirstElement(importStatement.getImportElements()); + if (firstImportElement == null) { + return ImportPriority.PROJECT; + } + resolved = firstImportElement.resolve(); } - ImportPriority relativeToPriority = source == null || containingFile == null - ? ImportPriority.BUILTIN - : getImportPriority(element, containingFile); - final int rc = priority.compareTo(relativeToPriority); - if (rc < 0) { - return true; + if (resolved == null) { + return ImportPriority.BUILTIN; } - if (rc == 0) { - return nameToImport.compareTo(relativeToName.toString()) < 0; + + final PsiFileSystemItem resolvedFileOrDir; + if (resolved instanceof PsiDirectory) { + resolvedFileOrDir = (PsiFileSystemItem)resolved; } - return false; + else { + resolvedFileOrDir = resolved.getContainingFile(); + } + return getImportPriority(importStatement, resolvedFileOrDir); + } + + @NotNull + public static ImportPriority getImportPriority(@NotNull PsiElement importLocation, @NotNull PsiFileSystemItem toImport) { + final VirtualFile vFile = toImport.getVirtualFile(); + if (vFile == null) { + return ImportPriority.PROJECT; + } + final ProjectRootManager projectRootManager = ProjectRootManager.getInstance(toImport.getProject()); + if (projectRootManager.getFileIndex().isInContent(vFile)) { + return ImportPriority.PROJECT; + } + final Module module = ModuleUtilCore.findModuleForPsiElement(importLocation); + final Sdk pythonSdk = module != null ? PythonSdkType.findPythonSdk(module) : projectRootManager.getProjectSdk(); + + return PythonSdkType.isStdLib(vFile, pythonSdk) ? ImportPriority.BUILTIN : ImportPriority.THIRD_PARTY; } /** @@ -214,7 +274,7 @@ public class AddImportHelper { insertParent.addAfter(importNodeToInsert, anchor); } else { - insertParent.addBefore(importNodeToInsert, getInsertPosition(insertParent, name, priority)); + insertParent.addBefore(importNodeToInsert, getInsertPosition(insertParent, importNodeToInsert, priority)); } } catch (IncorrectOperationException e) { @@ -254,7 +314,7 @@ public class AddImportHelper { insertParent = file; } if (InjectedLanguageManager.getInstance(file.getProject()).isInjectedFragment(file)) { - final PsiElement element = insertParent.addBefore(nodeToInsert, getInsertPosition(insertParent, from, priority)); + final PsiElement element = insertParent.addBefore(nodeToInsert, getInsertPosition(insertParent, nodeToInsert, priority)); PsiElement whitespace = element.getNextSibling(); if (!(whitespace instanceof PsiWhiteSpace)) { whitespace = PsiParserFacade.SERVICE.getInstance(file.getProject()).createWhiteSpaceFromText(" >>> "); @@ -266,7 +326,7 @@ public class AddImportHelper { insertParent.addAfter(nodeToInsert, anchor); } else { - insertParent.addBefore(nodeToInsert, getInsertPosition(insertParent, from, priority)); + insertParent.addBefore(nodeToInsert, getInsertPosition(insertParent, nodeToInsert, priority)); } } } @@ -359,19 +419,4 @@ public class AddImportHelper { } } } - - public static ImportPriority getImportPriority(PsiElement importLocation, @NotNull PsiFileSystemItem toImport) { - final VirtualFile vFile = toImport.getVirtualFile(); - if (vFile == null) { - return ImportPriority.PROJECT; - } - final ProjectRootManager projectRootManager = ProjectRootManager.getInstance(toImport.getProject()); - if (projectRootManager.getFileIndex().isInContent(vFile)) { - return ImportPriority.PROJECT; - } - Module module = ModuleUtilCore.findModuleForPsiElement(importLocation); - Sdk pythonSdk = module != null ? PythonSdkType.findPythonSdk(module) : projectRootManager.getProjectSdk(); - - return PythonSdkType.isStdLib(vFile, pythonSdk) ? ImportPriority.BUILTIN : ImportPriority.THIRD_PARTY; - } } diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java index 2354b2345003..58bc2020ee95 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/PyImportOptimizer.java @@ -72,20 +72,11 @@ public class PyImportOptimizer implements ImportOptimizer { public static final Ordering BY_NAME_COMPARATOR = Ordering.natural() .nullsLast() - .onResultOf(new Function() { + .onResultOf(new Function() { + @Nullable @Override - public String apply(@Nullable PyImportStatementBase importStatement) { - QualifiedName qualifiedName = null; - if (importStatement instanceof PyFromImportStatement) { - qualifiedName = ((PyFromImportStatement)importStatement).getImportSourceQName(); - } - else if (importStatement instanceof PyImportStatement) { - final PyImportElement importElement = ArrayUtil.getFirstElement(importStatement.getImportElements()); - if (importElement != null) { - qualifiedName = importElement.getImportedQName(); - } - } - return qualifiedName == null ? null : qualifiedName.toString(); + public QualifiedName apply(@NotNull PyImportStatementBase importStatement) { + return AddImportHelper.getImportFirstQualifiedName(importStatement); } }); diff --git a/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java b/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java index 2548e938eda6..525d6a96da9b 100644 --- a/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java +++ b/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java @@ -85,8 +85,11 @@ public class PythonFormattingModelBuilder implements FormattingModelBuilderEx, C .between(STATEMENT_OR_DECLARATION, FUNCTION_DECLARATION).blankLines(commonSettings.BLANK_LINES_AROUND_METHOD) .after(FUNCTION_DECLARATION).blankLines(commonSettings.BLANK_LINES_AROUND_METHOD) .after(CLASS_DECLARATION).blankLines(commonSettings.BLANK_LINES_AROUND_CLASS) + // Remove excess blank lines between imports, because ImportOptimizer gets rid of them anyway. + // Empty lines between import groups are handles in PyBlock#getSpacing + .between(IMPORT_STATEMENTS, IMPORT_STATEMENTS).spacing(0, Integer.MAX_VALUE, 1, false, 0) .between(STATEMENT_OR_DECLARATION, STATEMENT_OR_DECLARATION).spacing(0, Integer.MAX_VALUE, 1, false, 1) - + .between(COLON, STATEMENT_LIST).spacing(1, Integer.MAX_VALUE, 0, true, 0) .afterInside(COLON, TokenSet.create(KEY_VALUE_EXPRESSION, LAMBDA_EXPRESSION)).spaceIf(pySettings.SPACE_AFTER_PY_COLON) diff --git a/python/testData/addImport/newFirstImportInBuiltinGroup/a.py b/python/testData/addImport/newFirstImportInBuiltinGroup/a.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/addImport/newFirstImportInBuiltinGroup/main.after.py b/python/testData/addImport/newFirstImportInBuiltinGroup/main.after.py new file mode 100644 index 000000000000..002b4d10ecf5 --- /dev/null +++ b/python/testData/addImport/newFirstImportInBuiltinGroup/main.after.py @@ -0,0 +1,5 @@ +import datetime +import sys +import a + +print(datetime, sys, a) \ No newline at end of file diff --git a/python/testData/addImport/newFirstImportInBuiltinGroup/main.py b/python/testData/addImport/newFirstImportInBuiltinGroup/main.py new file mode 100644 index 000000000000..d6d4a59aaae6 --- /dev/null +++ b/python/testData/addImport/newFirstImportInBuiltinGroup/main.py @@ -0,0 +1,4 @@ +import sys +import a + +print(datetime, sys, a) \ No newline at end of file diff --git a/python/testData/addImport/newFirstImportInProjectGroup/a.py b/python/testData/addImport/newFirstImportInProjectGroup/a.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/addImport/newFirstImportInProjectGroup/b.py b/python/testData/addImport/newFirstImportInProjectGroup/b.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/addImport/newFirstImportInProjectGroup/main.after.py b/python/testData/addImport/newFirstImportInProjectGroup/main.after.py new file mode 100644 index 000000000000..7b3c94c17fbf --- /dev/null +++ b/python/testData/addImport/newFirstImportInProjectGroup/main.after.py @@ -0,0 +1,6 @@ +import sys + +import a +import b + +print(sys, a, b) \ No newline at end of file diff --git a/python/testData/addImport/newFirstImportInProjectGroup/main.py b/python/testData/addImport/newFirstImportInProjectGroup/main.py new file mode 100644 index 000000000000..1db88cb4a9f0 --- /dev/null +++ b/python/testData/addImport/newFirstImportInProjectGroup/main.py @@ -0,0 +1,4 @@ +import sys +import b + +print(sys, a, b) \ No newline at end of file diff --git a/python/testData/addImport/newFirstImportInProjectGroupWithExistingBlankLineAbove/a.py b/python/testData/addImport/newFirstImportInProjectGroupWithExistingBlankLineAbove/a.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/addImport/newFirstImportInProjectGroupWithExistingBlankLineAbove/b.py b/python/testData/addImport/newFirstImportInProjectGroupWithExistingBlankLineAbove/b.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/addImport/newFirstImportInProjectGroupWithExistingBlankLineAbove/main.after.py b/python/testData/addImport/newFirstImportInProjectGroupWithExistingBlankLineAbove/main.after.py new file mode 100644 index 000000000000..7b3c94c17fbf --- /dev/null +++ b/python/testData/addImport/newFirstImportInProjectGroupWithExistingBlankLineAbove/main.after.py @@ -0,0 +1,6 @@ +import sys + +import a +import b + +print(sys, a, b) \ No newline at end of file diff --git a/python/testData/addImport/newFirstImportInProjectGroupWithExistingBlankLineAbove/main.py b/python/testData/addImport/newFirstImportInProjectGroupWithExistingBlankLineAbove/main.py new file mode 100644 index 000000000000..6b7ddec3032f --- /dev/null +++ b/python/testData/addImport/newFirstImportInProjectGroupWithExistingBlankLineAbove/main.py @@ -0,0 +1,5 @@ +import sys + +import b + +print(sys, a, b) \ No newline at end of file diff --git a/python/testData/addImport/newLastImportInBuiltinGroup/a.py b/python/testData/addImport/newLastImportInBuiltinGroup/a.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/addImport/newLastImportInBuiltinGroup/main.after.py b/python/testData/addImport/newLastImportInBuiltinGroup/main.after.py new file mode 100644 index 000000000000..7f4ff331f627 --- /dev/null +++ b/python/testData/addImport/newLastImportInBuiltinGroup/main.after.py @@ -0,0 +1,6 @@ +import datetime +import sys + +import a + +print(datetime, sys, a) \ No newline at end of file diff --git a/python/testData/addImport/newLastImportInBuiltinGroup/main.py b/python/testData/addImport/newLastImportInBuiltinGroup/main.py new file mode 100644 index 000000000000..d0e5409586f4 --- /dev/null +++ b/python/testData/addImport/newLastImportInBuiltinGroup/main.py @@ -0,0 +1,4 @@ +import datetime +import a + +print(datetime, sys, a) \ No newline at end of file diff --git a/python/testData/addImport/newLastImportInProjectGroup/a.py b/python/testData/addImport/newLastImportInProjectGroup/a.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/addImport/newLastImportInProjectGroup/main.after.py b/python/testData/addImport/newLastImportInProjectGroup/main.after.py new file mode 100644 index 000000000000..6b78646ab313 --- /dev/null +++ b/python/testData/addImport/newLastImportInProjectGroup/main.after.py @@ -0,0 +1,5 @@ +import sys +import a +import b + +print(sys, a, b) \ No newline at end of file diff --git a/python/testData/addImport/newLastImportInProjectGroup/main.py b/python/testData/addImport/newLastImportInProjectGroup/main.py new file mode 100644 index 000000000000..1a5843fae363 --- /dev/null +++ b/python/testData/addImport/newLastImportInProjectGroup/main.py @@ -0,0 +1,4 @@ +import sys +import a + +print(sys, a, b) \ No newline at end of file diff --git a/python/testData/addImport/newThirdPartyImportInBetween/a.py b/python/testData/addImport/newThirdPartyImportInBetween/a.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/addImport/newThirdPartyImportInBetween/main.after.py b/python/testData/addImport/newThirdPartyImportInBetween/main.after.py new file mode 100644 index 000000000000..145a3f4cb7a3 --- /dev/null +++ b/python/testData/addImport/newThirdPartyImportInBetween/main.after.py @@ -0,0 +1,7 @@ +import sys + +import third_party + +import a + +print(sys, third_party, a) \ No newline at end of file diff --git a/python/testData/addImport/newThirdPartyImportInBetween/main.py b/python/testData/addImport/newThirdPartyImportInBetween/main.py new file mode 100644 index 000000000000..d673ba37a81b --- /dev/null +++ b/python/testData/addImport/newThirdPartyImportInBetween/main.py @@ -0,0 +1,4 @@ +import sys +import a + +print(sys, third_party, a) \ No newline at end of file diff --git a/python/testData/inspections/MoveFromFutureImportDocString_after.py b/python/testData/inspections/MoveFromFutureImportDocString_after.py index 86ac79b6053e..5f8f6e87ee0c 100644 --- a/python/testData/inspections/MoveFromFutureImportDocString_after.py +++ b/python/testData/inspections/MoveFromFutureImportDocString_after.py @@ -1,6 +1,5 @@ """This is a docstring.""" from __future__ import with_statement - from __future__ import print_function #comment from __future__ import absolute_import diff --git a/python/testSrc/com/jetbrains/python/PyAddImportTest.java b/python/testSrc/com/jetbrains/python/PyAddImportTest.java index c0b2ee5bb214..c00c9f0ef076 100644 --- a/python/testSrc/com/jetbrains/python/PyAddImportTest.java +++ b/python/testSrc/com/jetbrains/python/PyAddImportTest.java @@ -19,44 +19,58 @@ 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.codeInsight.imports.AddImportHelper.ImportPriority; 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; +import static com.jetbrains.python.codeInsight.imports.AddImportHelper.ImportPriority.*; + /** * @author yole */ public class PyAddImportTest extends PyTestCase { 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, null); - } - }.execute(); - myFixture.checkResultByFile("addImport/addBuiltin.after.py"); + doAddImport("re", BUILTIN); } public void testParens() { // PY-7400 - doAddImportFrom("urllib", "unquote_plus"); + doAddImportFrom("urllib", "unquote_plus", BUILTIN); } public void testComment() { // PY-8034 - doAddImportFrom("urllib", "unquote_plus"); + doAddImportFrom("urllib", "unquote_plus", BUILTIN); } - private void doAddImportFrom(final String path, final String name) { - myFixture.configureByFile("addImport/" + getTestName(true) + ".py"); - new WriteCommandAction(myFixture.getProject(), myFixture.getFile()) { - @Override - protected void run(Result result) throws Throwable { - AddImportHelper.addOrUpdateFromImportStatement(myFixture.getFile(), path, name, null, AddImportHelper.ImportPriority.BUILTIN, null); - } - }.execute(); - myFixture.checkResultByFile("addImport/" + getTestName(true) + ".after.py"); + // PY-14765 + public void testNewFirstImportInBuiltinGroup() { + doAddImportWithResolveInProject("datetime", BUILTIN); + } + + // PY-14765 + public void testNewLastImportInBuiltinGroup() { + doAddImportWithResolveInProject("sys", BUILTIN); + } + + // PY-14765 + public void testNewFirstImportInProjectGroup() { + doAddImportWithResolveInProject("a", PROJECT); + } + + // PY-14765 + public void testNewFirstImportInProjectGroupWithExistingBlankLineAbove() { + doAddImportWithResolveInProject("a", PROJECT); + } + + // PY-14765 + public void testNewLastImportInProjectGroup() { + doAddImportWithResolveInProject("b", PROJECT); + } + + // PY-14765 + public void testNewThirdPartyImportInBetween() { + doAddImportWithResolveInProject("third_party", THIRD_PARTY); } // PY-6020 @@ -79,25 +93,66 @@ public class PyAddImportTest extends PyTestCase { testLocalImport(); } + private void doAddImportFrom(final String path, final String name, final ImportPriority priority) { + myFixture.configureByFile(getTestName(true) + ".py"); + WriteCommandAction.runWriteCommandAction(myFixture.getProject(), new Runnable() { + @Override + public void run() { + AddImportHelper.addOrUpdateFromImportStatement(myFixture.getFile(), path, name, null, priority, null); + } + }); + myFixture.checkResultByFile(getTestName(true) + ".after.py"); + } + + private void doAddImport(final String name, final ImportPriority priority) { + myFixture.configureByFile(getTestName(true) + ".py"); + WriteCommandAction.runWriteCommandAction(myFixture.getProject(), new Runnable() { + @Override + public void run() { + AddImportHelper.addImportStatement(myFixture.getFile(), name, null, priority, null); + } + }); + myFixture.checkResultByFile(getTestName(true) + ".after.py"); + } + + private void doAddImportWithResolveInProject(final String name, final ImportPriority priority) { + final String testName = getTestName(true); + myFixture.copyDirectoryToProject(testName, ""); + myFixture.configureByFile("main.py"); + WriteCommandAction.runWriteCommandAction(myFixture.getProject(), new Runnable() { + @Override + public void run() { + AddImportHelper.addImportStatement(myFixture.getFile(), name, null, priority, null); + } + }); + myFixture.checkResultByFile(testName + "/main.after.py"); + } + /** * 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"); + myFixture.configureByFile(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); + AddImportHelper.addLocalFromImportStatement(reference.getElement(), qualifier, name); } else { - AddImportHelper.addLocalImportStatement((PyElement)reference.getElement(), name); + AddImportHelper.addLocalImportStatement(reference.getElement(), name); } } }.execute(); - myFixture.checkResultByFile("addImport/" + getTestName(true) + ".after.py"); + myFixture.checkResultByFile(getTestName(true) + ".after.py"); + } + + @Override + protected String getTestDataPath() { + return super.getTestDataPath() + "/addImport"; } }