PY-46358 Add the actual "as" part to local imports for common package aliases

GitOrigin-RevId: ad1061472e98c66ce3abc0e6d5e633ab271c2b41
This commit is contained in:
Mikhail Golubev
2021-01-07 17:25:43 +00:00
committed by intellij-monorepo-bot
parent 6f5997ee45
commit 4ca188384d
11 changed files with 41 additions and 37 deletions
@@ -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);
}
}
@@ -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
@@ -0,0 +1,2 @@
def main():
<error descr="Unresolved reference 'plt'">p<caret>lt</error>.plot
@@ -0,0 +1,3 @@
def main():
from matplotlib import pyplot as plt
plt.plot
@@ -0,0 +1,2 @@
def main():
<error descr="Unresolved reference 'np'">n<caret>p</error>.ndarray
@@ -0,0 +1,3 @@
def main():
import numpy as np
np.ndarray
@@ -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");
@@ -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<String> candidates = ContainerUtil.map(fix.getCandidates(), c -> c.getPresentableText());