diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java b/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java index 11a3dd090178..3fbeb3add142 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java +++ b/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java @@ -90,7 +90,7 @@ public abstract class PyElementGenerator { @NotNull public abstract PyCallExpression createCallExpression(final LanguageLevel langLevel, String functionName); - public abstract PyImportElement createImportElement(final LanguageLevel languageLevel, String name); + public abstract PyImportElement createImportElement(@NotNull LanguageLevel languageLevel, @NotNull String name, @Nullable String alias); public abstract PyFunction createProperty(final LanguageLevel languageLevel, String propertyName, diff --git a/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java b/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java index 45b9d40e8825..c58972997246 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java @@ -462,7 +462,7 @@ public class AddImportHelper { } } final PyElementGenerator generator = PyElementGenerator.getInstance(file.getProject()); - final PyImportElement importElement = generator.createImportElement(LanguageLevel.forElement(file), name); + final PyImportElement importElement = generator.createImportElement(LanguageLevel.forElement(file), name, asName); existingImport.add(importElement); return false; } diff --git a/python/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java b/python/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java index 02b95124f3dd..2c7050de44f8 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java @@ -200,7 +200,7 @@ public class ImportFromExistingAction implements QuestionAction { PsiElement parent = src.getParent(); if (parent instanceof PyFromImportStatement) { // add another import element right after the one we got - PsiElement newImportElement = gen.createImportElement(LanguageLevel.getDefault(), myName); + PsiElement newImportElement = gen.createImportElement(LanguageLevel.getDefault(), myName, null); parent.add(newImportElement); } else { // just 'import' diff --git a/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java b/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java index 336db0a31f1c..3e867147cf1b 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java @@ -299,8 +299,9 @@ public class PyElementGeneratorImpl extends PyElementGenerator { } @Override - public PyImportElement createImportElement(final LanguageLevel languageLevel, String name) { - return createFromText(languageLevel, PyImportElement.class, "from foo import " + name, new int[]{0, 6}); + public PyImportElement createImportElement(@NotNull final LanguageLevel languageLevel, @NotNull String name, @Nullable String alias) { + final String importStatement = "from foo import " + name + (alias != null ? " as " + alias : ""); + return createFromText(languageLevel, PyImportElement.class, importStatement, new int[]{0, 6}); } @Override diff --git a/python/testData/refactoring/move/fromImportAliases/after/src/a.py b/python/testData/refactoring/move/fromImportAliases/after/src/a.py new file mode 100644 index 000000000000..b28b04f64312 --- /dev/null +++ b/python/testData/refactoring/move/fromImportAliases/after/src/a.py @@ -0,0 +1,3 @@ + + + diff --git a/python/testData/refactoring/move/fromImportAliases/after/src/b.py b/python/testData/refactoring/move/fromImportAliases/after/src/b.py new file mode 100644 index 000000000000..4cdbb56b80d4 --- /dev/null +++ b/python/testData/refactoring/move/fromImportAliases/after/src/b.py @@ -0,0 +1,6 @@ +from lib import function1 as f1, function2 as f2 + + +def func(): + print(f1(), f2()) + return f2() \ No newline at end of file diff --git a/python/testData/refactoring/move/fromImportAliases/after/src/lib.py b/python/testData/refactoring/move/fromImportAliases/after/src/lib.py new file mode 100644 index 000000000000..279b5630edd8 --- /dev/null +++ b/python/testData/refactoring/move/fromImportAliases/after/src/lib.py @@ -0,0 +1,6 @@ +def function1(): + pass + + +def function2(): + pass diff --git a/python/testData/refactoring/move/fromImportAliases/before/src/a.py b/python/testData/refactoring/move/fromImportAliases/before/src/a.py new file mode 100644 index 000000000000..c2f0d99d4c38 --- /dev/null +++ b/python/testData/refactoring/move/fromImportAliases/before/src/a.py @@ -0,0 +1,6 @@ +from lib import function1 as f1, function2 as f2 + + +def func(): + print(f1(), f2()) + return f2() diff --git a/python/testData/refactoring/move/fromImportAliases/before/src/lib.py b/python/testData/refactoring/move/fromImportAliases/before/src/lib.py new file mode 100644 index 000000000000..279b5630edd8 --- /dev/null +++ b/python/testData/refactoring/move/fromImportAliases/before/src/lib.py @@ -0,0 +1,6 @@ +def function1(): + pass + + +def function2(): + pass diff --git a/python/testSrc/com/jetbrains/python/refactoring/PyMoveTest.java b/python/testSrc/com/jetbrains/python/refactoring/PyMoveTest.java index f336bb5e99d2..21a0c0de7aed 100644 --- a/python/testSrc/com/jetbrains/python/refactoring/PyMoveTest.java +++ b/python/testSrc/com/jetbrains/python/refactoring/PyMoveTest.java @@ -356,6 +356,11 @@ public class PyMoveTest extends PyTestCase { doMoveSymbolTest("use_foo", "c.py"); } + // PY-21366 + public void testFromImportAliases() { + doMoveSymbolTest("func", "b.py"); + } + private void doMoveFileTest(String fileName, String toDirName) { Project project = myFixture.getProject(); PsiManager manager = PsiManager.getInstance(project);