PY-21366 Consider import aliases when restoring imports for moved element

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,

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;
       }

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'

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

python/testData/refactoring/move/fromImportAliases/after/src/a.py

@@ -0,0 +1,3 @@
+
+
+

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()

python/testData/refactoring/move/fromImportAliases/after/src/lib.py

@@ -0,0 +1,6 @@
+def function1():
+    pass
+
+
+def function2():
+    pass

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()

python/testData/refactoring/move/fromImportAliases/before/src/lib.py

@@ -0,0 +1,6 @@
+def function1():
+    pass
+
+
+def function2():
+    pass

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);