diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java index 34f09828d191..66e8c8935e8c 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyFromImportStatementImpl.java @@ -187,32 +187,41 @@ public class PyFromImportStatementImpl extends PyBaseElementImpla, b, c) diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImport/main_after.py b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportBeginning/main_after.py similarity index 100% rename from python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImport/main_after.py rename to python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportBeginning/main_after.py diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportEnd/lib.py b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportEnd/lib.py new file mode 100644 index 000000000000..1e006da2fe73 --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportEnd/lib.py @@ -0,0 +1 @@ +a, b, c = 1, 2, 3 diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportEnd/main.py b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportEnd/main.py new file mode 100644 index 000000000000..2a364ca89be2 --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportEnd/main.py @@ -0,0 +1,3 @@ +from lib import a, b + +print(a, b, c) diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportEnd/main_after.py b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportEnd/main_after.py new file mode 100644 index 000000000000..6fe9292ff21f --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportEnd/main_after.py @@ -0,0 +1,3 @@ +from lib import a, b, c + +print(a, b, c) diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportInTheMiddle/lib.py b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportInTheMiddle/lib.py new file mode 100644 index 000000000000..1e006da2fe73 --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportInTheMiddle/lib.py @@ -0,0 +1 @@ +a, b, c = 1, 2, 3 diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImport/main.py b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportInTheMiddle/main.py similarity index 100% rename from python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImport/main.py rename to python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportInTheMiddle/main.py diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportInTheMiddle/main_after.py b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportInTheMiddle/main_after.py new file mode 100644 index 000000000000..6fe9292ff21f --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/orderingOfNamesInFromImportInTheMiddle/main_after.py @@ -0,0 +1,3 @@ +from lib import a, b, c + +print(a, b, c) diff --git a/python/testSrc/com/jetbrains/python/PyMiscellaneousPsiOperationsTest.java b/python/testSrc/com/jetbrains/python/PyMiscellaneousPsiOperationsTest.java index fa869e766d4d..cd2b605d325c 100644 --- a/python/testSrc/com/jetbrains/python/PyMiscellaneousPsiOperationsTest.java +++ b/python/testSrc/com/jetbrains/python/PyMiscellaneousPsiOperationsTest.java @@ -16,10 +16,9 @@ package com.jetbrains.python; import com.intellij.psi.util.QualifiedName; +import com.intellij.util.containers.ContainerUtil; import com.jetbrains.python.fixtures.PyTestCase; -import com.jetbrains.python.psi.LanguageLevel; -import com.jetbrains.python.psi.PyElementGenerator; -import com.jetbrains.python.psi.PyQualifiedExpression; +import com.jetbrains.python.psi.*; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -41,6 +40,54 @@ public class PyMiscellaneousPsiOperationsTest extends PyTestCase { checkAsQualifiedNameResult("-foo + bar", null); } + public void testAddingNameInFromImportStatement() { + checkAddingNameInFromImport("from mod import foo", "bar", "foo", true, "from mod import bar, foo"); + checkAddingNameInFromImport("from mod import foo", "bar", "foo", false, "from mod import foo, bar"); + checkAddingNameInFromImport("from mod import foo", "bar", null, false, "from mod import bar, foo"); + checkAddingNameInFromImport("from mod import foo", "bar", null, true, "from mod import foo, bar"); + checkAddingNameInFromImport("from mod import (foo) # comment", "bar", "foo", true, "from mod import (bar, foo) # comment"); + checkAddingNameInFromImport("from mod import (foo) # comment", "bar", "foo", false, "from mod import (foo, bar) # comment"); + checkAddingNameInFromImport("from mod import (foo,)", "bar", "foo", false, "from mod import (foo, bar,)"); + checkAddingNameInFromImport("from mod import (foo # comment\n" + + " )", + "bar", "foo", false, + "from mod import (foo, bar # comment\n" + + " )"); + checkAddingNameInFromImport("from mod import ", "bar", null, false, "from mod import bar"); + checkAddingNameInFromImport("from mod import ", "bar", null, true, "from mod import bar"); + checkAddingNameInFromImport("from mod import (", "bar", null, true, "from mod import (bar"); + // TODO change where the placeholder empty import element is added in such cases + //checkAddingNameInFromImport("from mod import ( # comment", "bar", null, true, "from mod import (bar # comment"); + checkAddingNameInFromImport("from mod import ()", "bar", null, true, "from mod import (bar)"); + } + + private void checkAddingNameInFromImport(@NotNull String fromImport, + @NotNull String newName, + @Nullable String anchorName, + boolean before, + @NotNull String result) { + final PyElementGenerator generator = PyElementGenerator.getInstance(myFixture.getProject()); + final LanguageLevel languageLevel = LanguageLevel.PYTHON27; + final PyFromImportStatement fromImportElem = generator.createFromText(languageLevel, PyFromImportStatement.class, fromImport, new int[]{0}); + final PyImportElement anchor; + if (anchorName != null) { + anchor = ContainerUtil.find(fromImportElem.getImportElements(), + importElem -> importElem.getImportedQName().toString().equals(anchorName)); + assertNotNull(anchor); + } + else { + anchor = null; + } + final PyImportElement newNameElem = generator.createImportElement(languageLevel, newName, null); + if (before) { + fromImportElem.addBefore(newNameElem, anchor); + } + else { + fromImportElem.addAfter(newNameElem, anchor); + } + assertEquals(result, fromImportElem.getText()); + } + private void checkAsQualifiedNameResult(@NotNull String expression, @Nullable QualifiedName expectedQualifiedName) { final PyElementGenerator generator = PyElementGenerator.getInstance(myFixture.getProject()); final PyQualifiedExpression expr = (PyQualifiedExpression)generator.createExpressionFromText(LanguageLevel.PYTHON27, expression); diff --git a/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java b/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java index 80cf978ac951..a58d57fa7587 100644 --- a/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java +++ b/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java @@ -224,7 +224,21 @@ public class PyAddImportQuickFixTest extends PyQuickFixTestCase { } // PY-23968 - public void testOrderingOfNamesInFromImport() { + public void testOrderingOfNamesInFromImportBeginning() { + getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_IMPORTS = true; + getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_NAMES_IN_FROM_IMPORTS = true; + doMultiFileAutoImportTest("Import"); + } + + // PY-23968 + public void testOrderingOfNamesInFromImportInTheMiddle() { + getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_IMPORTS = true; + getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_NAMES_IN_FROM_IMPORTS = true; + doMultiFileAutoImportTest("Import"); + } + + // PY-23968 + public void testOrderingOfNamesInFromImportEnd() { getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_IMPORTS = true; getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_NAMES_IN_FROM_IMPORTS = true; doMultiFileAutoImportTest("Import");