diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java index a0c39523342a..2f3a1afaeee6 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java @@ -1,9 +1,6 @@ // Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. package com.jetbrains.python.codeInsight.imports; -import static com.jetbrains.python.psi.PyUtil.as; -import static com.jetbrains.python.psi.PyUtil.sure; - import com.intellij.lang.injection.InjectedLanguageManager; import com.intellij.openapi.diagnostic.Logger; import com.intellij.openapi.module.Module; @@ -14,16 +11,7 @@ import com.intellij.openapi.roots.ProjectRootManager; import com.intellij.openapi.util.registry.Registry; import com.intellij.openapi.util.text.StringUtil; import com.intellij.openapi.vfs.VirtualFile; -import com.intellij.psi.PsiComment; -import com.intellij.psi.PsiDirectory; -import com.intellij.psi.PsiDirectoryContainer; -import com.intellij.psi.PsiElement; -import com.intellij.psi.PsiFile; -import com.intellij.psi.PsiFileSystemItem; -import com.intellij.psi.PsiLanguageInjectionHost; -import com.intellij.psi.PsiNamedElement; -import com.intellij.psi.PsiParserFacade; -import com.intellij.psi.PsiWhiteSpace; +import com.intellij.psi.*; import com.intellij.psi.codeStyle.CodeStyleManager; import com.intellij.psi.templateLanguages.OuterLanguageElement; import com.intellij.psi.util.PsiTreeUtil; @@ -36,29 +24,22 @@ import com.jetbrains.python.PythonCodeStyleService; import com.jetbrains.python.codeInsight.PyCodeInsightSettings; import com.jetbrains.python.documentation.docstrings.DocStringUtil; import com.jetbrains.python.documentation.doctest.PyDocstringFile; -import com.jetbrains.python.psi.LanguageLevel; -import com.jetbrains.python.psi.PyAssignmentStatement; -import com.jetbrains.python.psi.PyElement; -import com.jetbrains.python.psi.PyElementGenerator; -import com.jetbrains.python.psi.PyExpression; -import com.jetbrains.python.psi.PyFile; -import com.jetbrains.python.psi.PyFromImportStatement; -import com.jetbrains.python.psi.PyImportElement; -import com.jetbrains.python.psi.PyImportStatement; -import com.jetbrains.python.psi.PyImportStatementBase; -import com.jetbrains.python.psi.PyStatement; -import com.jetbrains.python.psi.PyUtil; +import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.resolve.QualifiedNameFinder; import com.jetbrains.python.pyi.PyiFile; import com.jetbrains.python.pyi.PyiUtil; import com.jetbrains.python.sdk.PythonSdkUtil; +import one.util.streamex.StreamEx; +import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; + import java.util.ArrayList; import java.util.Comparator; import java.util.List; import java.util.Objects; -import one.util.streamex.StreamEx; -import org.jetbrains.annotations.NotNull; -import org.jetbrains.annotations.Nullable; + +import static com.jetbrains.python.psi.PyUtil.as; +import static com.jetbrains.python.psi.PyUtil.sure; /** * Does the actual job of adding an import statement into a file. @@ -162,25 +143,28 @@ public final class AddImportHelper { private AddImportHelper() { } - public static void addLocalImportStatement(@NotNull PsiElement element, @NotNull String name) { + public static void addLocalImportStatement(@NotNull PsiElement element, @NotNull String name, @Nullable String asName) { final PyElementGenerator generator = PyElementGenerator.getInstance(element.getProject()); final LanguageLevel languageLevel = LanguageLevel.forElement(element); final PsiElement anchor = getLocalInsertPosition(element); final PsiElement parentElement = sure(anchor).getParent(); if (parentElement != null) { - parentElement.addBefore(generator.createImportStatement(languageLevel, name, null), anchor); + parentElement.addBefore(generator.createImportStatement(languageLevel, name, asName), anchor); } } - public static void addLocalFromImportStatement(@NotNull PsiElement element, @NotNull String qualifier, @NotNull String name) { + public static void addLocalFromImportStatement(@NotNull PsiElement element, + @NotNull String qualifier, + @NotNull String name, + @Nullable String asName) { final PyElementGenerator generator = PyElementGenerator.getInstance(element.getProject()); final LanguageLevel languageLevel = LanguageLevel.forElement(element); final PsiElement anchor = getLocalInsertPosition(element); final PsiElement parentElement = sure(anchor).getParent(); if (parentElement != null) { - parentElement.addBefore(generator.createFromImportStatement(languageLevel, qualifier, name, null), anchor); + parentElement.addBefore(generator.createFromImportStatement(languageLevel, qualifier, name, asName), anchor); } } diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java index 743b7481f1e6..b68afdc46931 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java @@ -124,7 +124,7 @@ public class ImportFromExistingAction implements QuestionAction { // A root-level module or package cannot be imported with a "from" import. if (PyUtil.isRoot(item.getFile())) { if (myImportLocally) { - AddImportHelper.addLocalImportStatement(myTarget, item.getImportableName()); + AddImportHelper.addLocalImportStatement(myTarget, item.getImportableName(), item.getAsName()); } else { AddImportHelper.addImportStatement(file, item.getImportableName(), item.getAsName(), priority, myTarget); @@ -138,7 +138,7 @@ public class ImportFromExistingAction implements QuestionAction { nameToImport += "." + item.getImportableName(); } if (myImportLocally) { - AddImportHelper.addLocalImportStatement(myTarget, nameToImport); + AddImportHelper.addLocalImportStatement(myTarget, nameToImport, item.getAsName()); } else { AddImportHelper.addImportStatement(file, nameToImport, item.getAsName(), priority, myTarget); @@ -149,7 +149,7 @@ public class ImportFromExistingAction implements QuestionAction { } else { if (myImportLocally) { - AddImportHelper.addLocalFromImportStatement(myTarget, qualifiedName, item.getImportableName()); + AddImportHelper.addLocalFromImportStatement(myTarget, qualifiedName, item.getImportableName(), item.getAsName()); } else { // "Update" scenario takes place inside injected fragments, for normal AST addToExistingImport() will be used instead diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/localFromImportForCommonPackageAlias/main.py b/python/testData/quickFixes/PyAddImportQuickFixTest/localFromImportForCommonPackageAlias/main.py new file mode 100644 index 000000000000..3a65b727d594 --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/localFromImportForCommonPackageAlias/main.py @@ -0,0 +1,2 @@ +def main(): + plt.plot diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/localFromImportForCommonPackageAlias/main_after.py b/python/testData/quickFixes/PyAddImportQuickFixTest/localFromImportForCommonPackageAlias/main_after.py new file mode 100644 index 000000000000..536a2b912389 --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/localFromImportForCommonPackageAlias/main_after.py @@ -0,0 +1,3 @@ +def main(): + from matplotlib import pyplot as plt + plt.plot diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/localFromImportForCommonPackageAlias/matplotlib/__init__.py b/python/testData/quickFixes/PyAddImportQuickFixTest/localFromImportForCommonPackageAlias/matplotlib/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/localFromImportForCommonPackageAlias/matplotlib/pyplot.py b/python/testData/quickFixes/PyAddImportQuickFixTest/localFromImportForCommonPackageAlias/matplotlib/pyplot.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/localPlainImportForCommonPackageAlias/main.py b/python/testData/quickFixes/PyAddImportQuickFixTest/localPlainImportForCommonPackageAlias/main.py new file mode 100644 index 000000000000..fcdaa0937574 --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/localPlainImportForCommonPackageAlias/main.py @@ -0,0 +1,2 @@ +def main(): + np.ndarray diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/localPlainImportForCommonPackageAlias/main_after.py b/python/testData/quickFixes/PyAddImportQuickFixTest/localPlainImportForCommonPackageAlias/main_after.py new file mode 100644 index 000000000000..1437bd1f07b4 --- /dev/null +++ b/python/testData/quickFixes/PyAddImportQuickFixTest/localPlainImportForCommonPackageAlias/main_after.py @@ -0,0 +1,3 @@ +def main(): + import numpy as np + np.ndarray diff --git a/python/testData/quickFixes/PyAddImportQuickFixTest/localPlainImportForCommonPackageAlias/numpy/__init__.py b/python/testData/quickFixes/PyAddImportQuickFixTest/localPlainImportForCommonPackageAlias/numpy/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testSrc/com/jetbrains/python/PyAddImportTest.java b/python/testSrc/com/jetbrains/python/PyAddImportTest.java index 008ae3b692ff..05d527d79cdd 100644 --- a/python/testSrc/com/jetbrains/python/PyAddImportTest.java +++ b/python/testSrc/com/jetbrains/python/PyAddImportTest.java @@ -234,10 +234,10 @@ public class PyAddImportTest extends PyTestCase { WriteCommandAction.runWriteCommandAction(myFixture.getProject(), () -> { final PsiReference reference = PyResolveTestCase.findReferenceByMarker(myFixture.getFile()); if (qualifier != null) { - AddImportHelper.addLocalFromImportStatement(reference.getElement(), qualifier, name); + AddImportHelper.addLocalFromImportStatement(reference.getElement(), qualifier, name, null); } else { - AddImportHelper.addLocalImportStatement(reference.getElement(), name); + AddImportHelper.addLocalImportStatement(reference.getElement(), name, null); } }); myFixture.checkResultByFile(getTestName(true) + ".after.py"); diff --git a/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java b/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java index 818e98e63909..bd62c295ca23 100644 --- a/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java +++ b/python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java @@ -316,6 +316,16 @@ public class PyAddImportQuickFixTest extends PyQuickFixTestCase { doMultiFileAutoImportTest("Import"); } + // PY-46358 + public void testLocalPlainImportForCommonPackageAlias() { + doMultiFileAutoImportTest("Import 'numpy as np' locally"); + } + + // PY-46358 + public void testLocalFromImportForCommonPackageAlias() { + doMultiFileAutoImportTest("Import 'matplotlib.pyplot as plt' locally"); + } + private void doTestProposedImportsOrdering(String @NotNull ... expected) { doMultiFileAutoImportTest("Import", fix -> { final List candidates = ContainerUtil.map(fix.getCandidates(), c -> c.getPresentableText());