PY-44026 Use import priority when adding import

GitOrigin-RevId: 1e3a64b971a73283863e3aba1f8251d6673f62f9

python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyModuleNameCompletionContributor.kt

@@ -48,7 +48,8 @@ class PyModuleNameCompletionContributor : CompletionContributor() {
     }
     val commonAlias = PyPackageAliasesProvider.commonImportAliases[result.prefixMatcher.prefix]
     if (commonAlias != null) {
-      result.addElement(LookupElementBuilder.create(result.prefixMatcher.prefix).withTypeText(commonAlias).withInsertHandler(packageInsertHandler))
+      result.addElement(
+        LookupElementBuilder.create(result.prefixMatcher.prefix).withTypeText(commonAlias).withInsertHandler(packageInsertHandler))
       return
     }
     getCompletionVariants(parameters.position.parent, parameters.originalFile).asSequence()
@@ -67,27 +68,22 @@ class PyModuleNameCompletionContributor : CompletionContributor() {
     resolveQualifiedName(QualifiedName.fromComponents(), fromFoothold(file))
       .asSequence()
       .filterIsInstance<PsiDirectory>()
-      .forEach { fillCompletionVariantsFromDir(it, element, result, alreadyAddedNames) }
+      .forEach {
+        val initPy = it.findFile(PyNames.INIT_DOT_PY)
+        if (initPy is PyFile) {
+          val moduleType = PyModuleType(initPy)
+          val context = ProcessingContext()
+          context.put(PyType.CTX_NAMES, alreadyAddedNames)
+          val completionVariants = moduleType.getCompletionVariants("", element, context)
+          result.addAll(listOf(*completionVariants))
+        }
+        else {
+          result.addAll(PyModuleType.getSubModuleVariants(it, element, alreadyAddedNames))
+        }
+      }
     return result
   }
 
-  private fun fillCompletionVariantsFromDir(targetDir: PsiDirectory?, element: PsiElement,
-                                            result: ArrayList<Any>, alreadyAddedNames: HashSet<String>) {
-    if (targetDir != null) {
-      val initPy = targetDir.findFile(PyNames.INIT_DOT_PY)
-      if (initPy is PyFile) {
-        val moduleType = PyModuleType((initPy as PyFile?)!!)
-        val context = ProcessingContext()
-        context.put(PyType.CTX_NAMES, alreadyAddedNames)
-        val completionVariants = moduleType.getCompletionVariants("", element, context)
-        result.addAll(listOf(*completionVariants))
-      }
-      else {
-        result.addAll(PyModuleType.getSubModuleVariants(targetDir, element, alreadyAddedNames))
-      }
-    }
-  }
-
   private fun shouldDoCompletion(parameters: CompletionParameters): Boolean {
     val element = parameters.position
     val parent = element.parent

python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyUnresolvedModuleAttributeCompletionContributor.kt

@@ -82,7 +82,9 @@ class PyUnresolvedModuleAttributeCompletionContributor : CompletionContributor()
           val name = QualifiedName.fromDottedString(item.lookupString).removeLastComponent().toString()
           val commonAlias = PyPackageAliasesProvider.commonImportAliases[name]
           val nameToImport = commonAlias ?: name
-          AddImportHelper.addImportStatement(context.file, nameToImport, if (commonAlias != null) name else null, null, ref?.element as? PyElement)
+          AddImportHelper.addImportStatement(context.file, nameToImport, if (commonAlias != null) name else null,
+                                             AddImportHelper.getImportPriority(context.file, psiElement.containingFile),
+                                             ref?.element as? PyElement)
         }
       }
     }
@@ -134,14 +136,15 @@ class PyUnresolvedModuleAttributeCompletionContributor : CompletionContributor()
           result.restartCompletionOnAnyPrefixChange()
           return
         }
-        val scope = PySearchUtilBase.excludeSdkTestsScope(project)
+        val scope = PySearchUtilBase.defaultSuggestionScope(parameters.originalFile)
         PyQualifiedNameCompletionMatcher.processMatchingExportedNames(
           packageName.append(attribute), if (commonAlias != null) qualifiedName else null, parameters.originalFile, scope,
           Processor {
             ProgressManager.checkCanceled()
             if (suggestedQualifiedNames.add(it.qualifiedName.toString())) {
               resultMatchingCompleteReference.addElement(LookupElementBuilder
-                                                           .createWithSmartPointer(it.qualifiedNameWithUserTypedAlias.toString(), it.element)
+                                                           .createWithSmartPointer(it.qualifiedNameWithUserTypedAlias.toString(),
+                                                                                   it.element)
                                                            .withIcon(it.element.getIcon(0))
                                                            .withInsertHandler(getInsertHandler(it.element, parameters.position)))
             }

python/python-psi-impl/src/com/jetbrains/python/psi/stubs/PyQualifiedNameCompletionMatcher.java

@@ -114,7 +114,7 @@ public class PyQualifiedNameCompletionMatcher {
       canonicalImportPath = QualifiedNameFinder.findCanonicalImportPath(element, currentFile);
     }
     else {
-      canonicalImportPath = QualifiedNameFinder.canonizeQualifiedName(moduleQualifiedName, currentFile);
+      canonicalImportPath = QualifiedNameFinder.canonizeQualifiedName(element, moduleQualifiedName, currentFile);
     }
     return canonicalImportPath;
   }

python/testData/completion/notImportedQualifiedName/UseImportPriorityWhenAddingImport/main.after.py

@@ -0,0 +1,12 @@
+import subprocess
+import sys
+
+import django.conf
+import django.utils.encoding
+
+import matplotlib.pyplot as plt
+
+subprocess.Popen
+
+sys.argv
+plt.func()

python/testData/completion/notImportedQualifiedName/UseImportPriorityWhenAddingImport/main.py

@@ -0,0 +1,10 @@
+import subprocess
+import sys
+
+import django.conf
+import django.utils.encoding
+
+subprocess.Popen
+
+sys.argv
+plt.<caret>

python/testData/completion/notImportedQualifiedName/UseImportPriorityWhenAddingImport/matplotlib/pyplot/__init__.py

@@ -0,0 +1,2 @@
+def func():
+    pass

python/testSrc/com/jetbrains/python/PyNotImportedPackageNameCompletionTest.java

@@ -1,11 +1,7 @@
 // Copyright 2000-2020 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;
 
-import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.fixtures.PyTestCase;
-import com.jetbrains.python.psi.stubs.PyQualifiedNameCompletionMatcher.QualifiedNameMatcher;
-
-import java.util.List;
 
 public class PyNotImportedPackageNameCompletionTest extends PyTestCase {
   public void testDotAfterPackageName() {

python/testSrc/com/jetbrains/python/PyNotImportedQualifiedNameCompletionTest.java

@@ -73,6 +73,13 @@ public class PyNotImportedQualifiedNameCompletionTest extends PyTestCase {
     myFixture.checkResultByFile(getTestName(true) + "/main.after.py");
   }
 
+  public void testUseImportPriorityWhenAddingImport() {
+    myFixture.copyDirectoryToProject(getTestName(false), "");
+    myFixture.configureByFile("main.py");
+    myFixture.completeBasic();
+    myFixture.checkResultByFile(getTestName(true) + "/main.after.py");
+  }
+
   @Override
   protected String getTestDataPath() {
     return super.getTestDataPath() + "/completion/notImportedQualifiedName/";