PY-36374 When suggesting imports, don't shadow package aliases with namesake definitions

Previously, we suggested variants for these well-known aliases only if no
other symbols were found. In particular, having "np" name defined in the
sources of pandas prevented "import numpy as np" from being offered.

ImportCandidateHolder had to be refactored a bit to keep the name of
the corresponding symbol itself, because it used to be taken from the parent
AutoImportQuickFix, which now only holds the name of the reference, and, thus,
"numpy" ended up being imported as "import np as np".

Also, such imports are now suggested higher in the list than the rest.

GitOrigin-RevId: 75bed654431e81e6df40f1400ee4b02a4ab0dcca

python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/ImportCandidateHolder.java

@@ -35,6 +35,7 @@ public class ImportCandidateHolder implements Comparable<ImportCandidateHolder>
   @Nullable private final SmartPsiElementPointer<PyImportElement> myImportElement;
   @NotNull private final SmartPsiElementPointer<PsiFileSystemItem> myFile;
   @Nullable private final QualifiedName myPath;
+  private final String myImportableName;
   @Nullable private final String myAsName;
   private final int myRelevance;
 
@@ -55,12 +56,12 @@ public class ImportCandidateHolder implements Comparable<ImportCandidateHolder>
     SmartPointerManager pointerManager = SmartPointerManager.getInstance(importable.getProject());
     myFile = pointerManager.createSmartPsiElementPointer(file);
     myImportable = pointerManager.createSmartPsiElementPointer(importable);
+    myImportableName = importable instanceof PsiNamedElement ? PyUtil.getElementNameWithoutExtension(((PsiNamedElement)importable)) : null;
     myImportElement = importElement != null ? pointerManager.createSmartPsiElementPointer(importElement) : null;
     myPath = path;
     myAsName = asName;
-    String name = importable instanceof PsiNamedElement ? ((PsiNamedElement)importable).getName() : null;
-    myRelevance = PyCompletionUtilsKt.computeCompletionWeight(importable, name, path, null, false);
-    LOG.debug("Computed relevance for import item ", name, ": ", myRelevance);
+    myRelevance = PyCompletionUtilsKt.computeCompletionWeight(importable, myImportableName, myPath, null, false);
+    LOG.debug("Computed relevance for import item ", myImportableName, ": ", myRelevance);
     assert importElement != null || path != null; // one of these must be present
   }
 
@@ -74,6 +75,10 @@ public class ImportCandidateHolder implements Comparable<ImportCandidateHolder>
     return myImportable.getElement();
   }
 
+  public String getImportableName() {
+    return myImportableName;
+  }
+
   @Nullable
   public PyImportElement getImportElement() {
     return myImportElement != null ? myImportElement.getElement() : null;
@@ -156,6 +161,9 @@ public class ImportCandidateHolder implements Comparable<ImportCandidateHolder>
   public int compareTo(@NotNull ImportCandidateHolder other) {
     if (myImportElement != null && other.myImportElement == null) return -1;
     if (myImportElement == null && other.myImportElement != null) return 1;
+    if (myAsName != null && other.myAsName == null) return -1;
+    if (myAsName == null && other.myAsName != null) return 1;
+
     int comparedRelevance = Comparator
       .comparing(ImportCandidateHolder::getRelevance).reversed()
       .compare(this, other);

python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java

@@ -10,7 +10,6 @@ import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiFileSystemItem;
-import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.ObjectUtils;
 import com.jetbrains.python.PyPsiBundle;
 import com.jetbrains.python.psi.*;
@@ -18,6 +17,7 @@ import com.jetbrains.python.psi.impl.PyPsiUtils;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.List;
+import java.util.Objects;
 
 import static com.jetbrains.python.psi.PyUtil.as;
 
@@ -99,16 +99,15 @@ public class ImportFromExistingAction implements QuestionAction {
   }
 
   private void doIt(final ImportCandidateHolder item) {
-    PyImportElement src = item.getImportElement();
-    if (src != null) {
-      addToExistingImport(src);
+    if (item.getImportElement() != null) {
+      addToExistingImport(item);
     }
     else { // no existing import, add it then use it
       addImportStatement(item);
     }
   }
 
-  private void addImportStatement(ImportCandidateHolder item) {
+  private void addImportStatement(@NotNull ImportCandidateHolder item) {
     final Project project = myTarget.getProject();
     final PyElementGenerator gen = PyElementGenerator.getInstance(project);
 
@@ -125,19 +124,18 @@ public class ImportFromExistingAction implements QuestionAction {
     // We are trying to import top-level module or package which thus cannot be qualified
     if (PyUtil.isRoot(item.getFile())) {
       if (myImportLocally) {
-        AddImportHelper.addLocalImportStatement(myTarget, myName);
+        AddImportHelper.addLocalImportStatement(myTarget, item.getImportableName());
       }
       else {
-        AddImportHelper.addImportStatement(file, myName, item.getAsName(), priority, myTarget);
+        AddImportHelper.addImportStatement(file, item.getImportableName(), item.getAsName(), priority, myTarget);
       }
     }
     else {
-      final QualifiedName path = item.getPath();
-      final String qualifiedName = path != null ? path.toString() : "";
+      final String qualifiedName = Objects.toString(item.getPath(), "");
       if (myUseQualifiedImport) {
         String nameToImport = qualifiedName;
         if (item.getImportable() instanceof PsiFileSystemItem) {
-          nameToImport += "." + myName;
+          nameToImport += "." + item.getImportableName();
         }
         if (myImportLocally) {
           AddImportHelper.addLocalImportStatement(myTarget, nameToImport);
@@ -149,27 +147,29 @@ public class ImportFromExistingAction implements QuestionAction {
       }
       else {
         if (myImportLocally) {
-          AddImportHelper.addLocalFromImportStatement(myTarget, qualifiedName, myName);
+          AddImportHelper.addLocalFromImportStatement(myTarget, qualifiedName, item.getImportableName());
         }
         else {
           // "Update" scenario takes place inside injected fragments, for normal AST addToExistingImport() will be used instead
-          AddImportHelper.addOrUpdateFromImportStatement(file, qualifiedName, myName, item.getAsName(), priority, myTarget);
+          AddImportHelper.addOrUpdateFromImportStatement(file, qualifiedName, item.getImportableName(), item.getAsName(), priority, myTarget);
         }
       }
     }
   }
 
 
-  private void addToExistingImport(PyImportElement src) {
-    final PyElementGenerator gen = PyElementGenerator.getInstance(myTarget.getProject());
+  private void addToExistingImport(@NotNull ImportCandidateHolder item) {
+    PyImportElement importElement = item.getImportElement();
+    assert importElement != null;
     // did user choose 'import' or 'from import'?
-    PsiElement parent = src.getParent();
+    PsiElement parent = importElement.getParent();
     if (parent instanceof PyFromImportStatement) {
-      AddImportHelper.addNameToFromImportStatement((PyFromImportStatement)parent, myName, null);
+      AddImportHelper.addNameToFromImportStatement((PyFromImportStatement)parent, item.getImportableName(), item.getAsName());
     }
     else { // just 'import'
       // all we need is to qualify our target
-      myTarget.replace(gen.createExpressionFromText(LanguageLevel.forElement(myTarget), src.getVisibleName() + "." + myName));
+      PyElementGenerator gen = PyElementGenerator.getInstance(myTarget.getProject());
+      myTarget.replace(gen.createExpressionFromText(LanguageLevel.forElement(myTarget), importElement.getVisibleName() + "." + myName));
     }
   }
 

python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/PyImportCollector.java

@@ -11,6 +11,7 @@ import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
+import com.jetbrains.python.inspections.unresolvedReference.PyPackageAliasesProvider;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyFileImpl;
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
@@ -113,27 +114,32 @@ public class PyImportCollector {
     }
     symbols.addAll(PyVariableNameIndex.find(myRefText, project, scope));
     if (isPossibleModuleReference()) {
-      symbols.addAll(findImportableModules(project, scope));
+      symbols.addAll(findImportableModules(myRefText, project, scope));
+      String packageQName = PyPackageAliasesProvider.commonImportAliases.get(myRefText);
+      if (packageQName != null) {
+        symbols.addAll(findImportableModules(packageQName, project, scope));
+      }
     }
-    if (!symbols.isEmpty()) {
-      for (PsiElement symbol : symbols) {
-        if (isIndexableTopLevel(symbol)) { // we only want top-level symbols
-          PsiFileSystemItem srcfile =
-            symbol instanceof PsiFileSystemItem ? ((PsiFileSystemItem)symbol).getParent() : symbol.getContainingFile();
-          if (srcfile != null && isAcceptableForImport(existingImportFile, srcfile)) {
-            QualifiedName importPath = QualifiedNameFinder.findCanonicalImportPath(symbol, myNode);
-            if (importPath == null) {
-              continue;
-            }
-            if (symbol instanceof PsiFileSystemItem) {
-              importPath = importPath.removeTail(1);
-            }
-            final String symbolImportQName = importPath.append(myRefText).toString();
-            if (!seenCandidateNames.contains(symbolImportQName)) {
-              // a new, valid hit
-              fix.addImport(symbol, srcfile, importPath, myAlias);
-              seenCandidateNames.add(symbolImportQName);
-            }
+    for (PsiElement symbol : symbols) {
+      if (isIndexableTopLevel(symbol)) { // we only want top-level symbols
+        PsiFileSystemItem srcfile =
+          symbol instanceof PsiFileSystemItem ? ((PsiFileSystemItem)symbol).getParent() : symbol.getContainingFile();
+        if (srcfile != null && isAcceptableForImport(existingImportFile, srcfile)) {
+          QualifiedName importPath = QualifiedNameFinder.findCanonicalImportPath(symbol, myNode);
+          if (importPath == null) {
+            continue;
+          }
+          if (symbol instanceof PsiFileSystemItem) {
+            importPath = importPath.removeTail(1);
+          }
+          if (!(symbol instanceof PsiNamedElement)) {
+            continue;
+          }
+          String name = PyUtil.getElementNameWithoutExtension((PsiNamedElement)symbol);
+          final String symbolImportQName = importPath.append(name).toString();
+          if (seenCandidateNames.add(symbolImportQName)) {
+            String alias = name.equals(myRefText) ? null : myRefText;
+            fix.addImport(symbol, srcfile, importPath, alias);
           }
         }
       }
@@ -186,10 +192,10 @@ public class PyImportCollector {
     return true;
   }
 
-  private Collection<PsiElement> findImportableModules(Project project, GlobalSearchScope scope) {
-    List<PsiElement> result = new ArrayList<>();
+  private Collection<PsiFileSystemItem> findImportableModules(String name, Project project, GlobalSearchScope scope) {
+    List<PsiFileSystemItem> result = new ArrayList<>();
     // Add packages
-    FilenameIndex.processFilesByName(myRefText, true, item -> {
+    FilenameIndex.processFilesByName(name, true, item -> {
       ProgressManager.checkCanceled();
       final PsiDirectory candidatePackageDir = as(item, PsiDirectory.class);
       if (candidatePackageDir != null && candidatePackageDir.findFile(PyNames.INIT_DOT_PY) != null) {
@@ -198,7 +204,7 @@ public class PyImportCollector {
       return true;
     }, scope, project, null);
     // Add modules
-    FilenameIndex.processFilesByName(myRefText + ".py", false, true, item -> {
+    FilenameIndex.processFilesByName(name + ".py", false, true, item -> {
       ProgressManager.checkCanceled();
       if (PyUtil.isImportable(myNode.getContainingFile(), item)) {
         result.add(item);

python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/PythonImportUtils.java

@@ -4,11 +4,11 @@ package com.jetbrains.python.codeInsight.imports;
 
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.module.ModuleUtilCore;
-import com.intellij.psi.*;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiReference;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
-import com.jetbrains.python.inspections.unresolvedReference.PyPackageAliasesProvider;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.sdk.PythonSdkUtil;
 import org.jetbrains.annotations.Nullable;
@@ -35,14 +35,7 @@ public final class PythonImportUtils {
       return null;
     }
 
-    AutoImportQuickFix fix = addCandidates(node, reference, refText, null);
-    if (fix != null) return fix;
-    final String packageName = PyPackageAliasesProvider.commonImportAliases.get(refText);
-    if (packageName != null) {
-      fix = addCandidates(node, reference, packageName, refText);
-      if (fix != null) return fix;
-    }
-    return null;
+    return addCandidates(node, reference, refText, null);
   }
 
   @Nullable

python/testData/quickFixes/PyAddImportQuickFixTest/commonPackageAlias/main.py

@@ -0,0 +1 @@
+<error descr="Unresolved reference 'np'">n<caret>p</error>.ndarray

python/testData/quickFixes/PyAddImportQuickFixTest/commonPackageAlias/main_after.py

@@ -0,0 +1,3 @@
+import numpy as np
+
+np.ndarray

python/testData/quickFixes/PyAddImportQuickFixTest/globalDefinitionDoesNotShadowCommonPackageAliasVariant/main.py

@@ -0,0 +1 @@
+<error descr="Unresolved reference 'np'">n<caret>p</error>.ndarray

python/testData/quickFixes/PyAddImportQuickFixTest/globalDefinitionDoesNotShadowCommonPackageAliasVariant/main_after.py

@@ -0,0 +1,3 @@
+import numpy as np
+
+np.ndarray

python/testData/quickFixes/PyAddImportQuickFixTest/globalDefinitionDoesNotShadowCommonPackageAliasVariant/pandas/__init__.py

@@ -0,0 +1 @@
+np = 42

python/testSrc/com/jetbrains/python/quickFixes/PyAddImportQuickFixTest.java

@@ -17,9 +17,12 @@ package com.jetbrains.python.quickFixes;
 
 import com.intellij.lang.injection.InjectedLanguageManager;
 import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.PsiDirectory;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFileSystemItem;
 import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.ObjectUtils;
 import com.intellij.util.Processor;
 import com.intellij.util.containers.ContainerUtil;
@@ -31,12 +34,15 @@ import com.jetbrains.python.formatter.PyCodeStyleSettings;
 import com.jetbrains.python.inspections.unresolvedReference.PyUnresolvedReferencesInspection;
 import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.psi.PyReferenceExpression;
+import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.List;
 import java.util.function.Consumer;
 
+import static com.jetbrains.python.psi.PyUtil.as;
+
 /**
  * @author Mikhail Golubev
  */
@@ -266,6 +272,25 @@ public class PyAddImportQuickFixTest extends PyQuickFixTestCase {
     doMultiFileAutoImportTest("Import");
   }
 
+  // PY-36374
+  public void testGlobalDefinitionDoesNotShadowCommonPackageAliasVariant() {
+    doMultiFileAutoImportTest("Import", fix -> {
+      List<ImportCandidateHolder> candidates = fix.getCandidates();
+      ImportCandidateHolder importNumpyAsNpVariant = ContainerUtil.find(candidates, c -> {
+        PsiDirectory dir = as(c.getImportable(), PsiDirectory.class);
+        return dir != null && dir.getName().equals("numpy") && "np".equals(c.getAsName());
+      });
+      assertNotNull(importNumpyAsNpVariant);
+      List<String> candidateText = ContainerUtil.map(fix.getCandidates(), c -> c.getPresentableText("np"));
+      assertContainsInRelativeOrder(candidateText, "np", "pandas.np");
+      return true;
+    });
+  }
+
+  public void testCommonPackageAlias() {
+    doMultiFileAutoImportTest("Import");
+  }
+
   private void doTestProposedImportsOrdering(@NotNull String text, String @NotNull ... expected) {
     doMultiFileAutoImportTest("Import", fix -> {
       final List<String> candidates = ContainerUtil.map(fix.getCandidates(), c -> c.getPresentableText(text));