PY-44026 Add completion with auto-import for common package aliases

GitOrigin-RevId: 11d6f0df866ba3ae35cfdb1aed231a3e357bb4e7

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

@@ -6,6 +6,7 @@ import com.intellij.codeInsight.completion.CompletionContributor
 import com.intellij.codeInsight.completion.CompletionParameters
 import com.intellij.codeInsight.completion.CompletionResultSet
 import com.intellij.codeInsight.completion.InsertHandler
+import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.psi.MultiplePsiFilesPerDocumentFileViewProvider
 import com.intellij.psi.PsiDirectory
@@ -14,6 +15,7 @@ import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.QualifiedName
 import com.intellij.util.ProcessingContext
 import com.jetbrains.python.PyNames
+import com.jetbrains.python.inspections.unresolvedReference.PyPackageAliasesProvider
 import com.jetbrains.python.psi.PyFile
 import com.jetbrains.python.psi.PyImportStatementBase
 import com.jetbrains.python.psi.PyReferenceExpression
@@ -38,17 +40,23 @@ class PyModuleNameCompletionContributor : CompletionContributor() {
 
   fun doFillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet) {
     val autoPopupController = AutoPopupController.getInstance(parameters.originalFile.project)
+    val packageInsertHandler = InsertHandler<LookupElement> { context, _ ->
+      // add dot for PyUnresolvedModuleAttributeCompletionContributor to work
+      context.document.insertString(context.tailOffset, ".")
+      context.editor.caretModel.moveToOffset(context.tailOffset)
+      autoPopupController.autoPopupMemberLookup(context.editor, null)
+    }
+    val commonAlias = PyPackageAliasesProvider.commonImportAliases[result.prefixMatcher.prefix]
+    if (commonAlias != null) {
+      result.addElement(LookupElementBuilder.create(result.prefixMatcher.prefix).withTypeText(commonAlias).withInsertHandler(packageInsertHandler))
+      return
+    }
     getCompletionVariants(parameters.position.parent, parameters.originalFile).asSequence()
       .filterIsInstance<LookupElementBuilder>()
       .filter { result.prefixMatcher.prefixMatches(it.lookupString) }
       .filterNot { it.lookupString.startsWith('_') }
       .map {
-        it.withInsertHandler(InsertHandler { context, _ ->
-          // add dot for PyUnresolvedModuleAttributeCompletionContributor to work
-          context.document.insertString(context.tailOffset, ".")
-          context.editor.caretModel.moveToOffset(context.tailOffset)
-          autoPopupController.autoPopupMemberLookup(context.editor, null)
-        })
+        it.withInsertHandler(packageInsertHandler)
       }
       .forEach { result.addElement(it) }
   }

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

@@ -18,6 +18,7 @@ import com.intellij.util.Processor
 import com.jetbrains.python.PyTokenTypes
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil
 import com.jetbrains.python.codeInsight.imports.AddImportHelper
+import com.jetbrains.python.inspections.unresolvedReference.PyPackageAliasesProvider
 import com.jetbrains.python.psi.*
 import com.jetbrains.python.psi.resolve.PyResolveUtil
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder
@@ -78,14 +79,17 @@ class PyUnresolvedModuleAttributeCompletionContributor : CompletionContributor()
       WriteCommandAction.writeCommandAction(context.project, context.file).run<RuntimeException> {
         val psiElement = item.psiElement
         if (psiElement is PsiNamedElement && psiElement.containingFile != null) {
-          val nameToImport = QualifiedName.fromDottedString(item.lookupString).removeLastComponent().toString()
-          AddImportHelper.addImportStatement(context.file, nameToImport, null, null, ref?.element as? PyElement)
+          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)
         }
       }
     }
 
     private fun getInsertHandler(elementToInsert: PyElement, position: PsiElement): InsertHandler<LookupElement> {
-      return if (elementToInsert is PyFunction && position.parent?.parent !is PyDecorator) functionInsertHandler else importingInsertHandler
+      return if (elementToInsert is PyFunction && position.parent?.parent !is PyDecorator) functionInsertHandler
+      else importingInsertHandler
     }
   }
 
@@ -104,19 +108,23 @@ class PyUnresolvedModuleAttributeCompletionContributor : CompletionContributor()
 
         ProgressManager.checkCanceled()
         val qualifiedName = qualifier.append(attribute)
+        val commonAlias = PyPackageAliasesProvider.commonImportAliases[qualifier.toString()]
+        val packageName = if (commonAlias != null) QualifiedName.fromDottedString(commonAlias) else qualifier
         val resultMatchingCompleteReference = result.withPrefixMatcher(QualifiedNameMatcher(qualifiedName))
-        PyModuleNameIndex.find(qualifier.lastComponent!!, project, true).asSequence()
-          .filter { QualifiedNameFinder.findShortestImportableQName(it) == qualifier }
+        PyModuleNameIndex.find(packageName.lastComponent!!, project, true).asSequence()
+          .filter { QualifiedNameFinder.findShortestImportableQName(it) == packageName }
           .flatMap { it.iterateNames().asSequence() }
+          .filter { it.containingFile != null }
           .filterNot { it is PsiFileSystemItem }
           .filterNot { it.name == null || it.name!!.startsWith('_') }
-          .filter { attribute.isEmpty() || result.prefixMatcher.prefixMatches(it.name!!) }
+          .filter { attribute.isEmpty() || resultMatchingCompleteReference.prefixMatcher.prefixMatches(it.name!!) }
           .mapNotNull {
             val qualifiedNameToSuggest = "$qualifier.${it.name}"
             if (suggestedQualifiedNames.add(qualifiedNameToSuggest)) {
               LookupElementBuilder.create(it, qualifiedNameToSuggest)
                 .withIcon(it.getIcon(0))
                 .withInsertHandler(getInsertHandler(it, parameters.position))
+                .withTypeText(commonAlias)
             }
             else null
           }
@@ -128,12 +136,12 @@ class PyUnresolvedModuleAttributeCompletionContributor : CompletionContributor()
         }
         val scope = PySearchUtilBase.excludeSdkTestsScope(project)
         PyQualifiedNameCompletionMatcher.processMatchingExportedNames(
-          qualifiedName, parameters.originalFile, scope,
+          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.qualifiedName.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

@@ -37,6 +37,7 @@ public class PyQualifiedNameCompletionMatcher {
   }
 
   public static void processMatchingExportedNames(@NotNull QualifiedName qualifiedNamePattern,
+                                                  @Nullable QualifiedName originallyTypedAlias,
                                                   @NotNull PsiFile currentFile,
                                                   @NotNull GlobalSearchScope scope,
                                                   @NotNull Processor<? super ExportedName> processor) {
@@ -71,7 +72,7 @@ public class PyQualifiedNameCompletionMatcher {
             }
             QualifiedName attributeQualifiedName = importPath.append(attributeName);
             if (alreadySuggestedAttributes.add(attributeQualifiedName)) {
-              if (!processor.process(new ExportedName(attributeQualifiedName, element))) {
+              if (!processor.process(new ExportedName(attributeQualifiedName, originallyTypedAlias, element))) {
                 return false;
               }
             }
@@ -120,10 +121,12 @@ public class PyQualifiedNameCompletionMatcher {
 
   public static final class ExportedName {
     private final QualifiedName myQualifiedName;
+    private final QualifiedName myOriginallyTypedQName;
     private final PyElement myElement;
 
-    private ExportedName(@NotNull QualifiedName qualifiedName, @NotNull PyElement element) {
+    private ExportedName(@NotNull QualifiedName qualifiedName, @Nullable QualifiedName originallyTypedQName, @NotNull PyElement element) {
       myQualifiedName = qualifiedName;
+      myOriginallyTypedQName = originallyTypedQName;
       myElement = element;
     }
 
@@ -132,10 +135,22 @@ public class PyQualifiedNameCompletionMatcher {
       return myQualifiedName;
     }
 
+    @Nullable
+    public QualifiedName getOriginallyTypedQName() {
+      return myOriginallyTypedQName;
+    }
+
     @NotNull
     public PyElement getElement() {
       return myElement;
     }
+
+    @NotNull
+    public QualifiedName getQualifiedNameWithUserTypedAlias() {
+      return myOriginallyTypedQName != null
+             ? myOriginallyTypedQName.removeLastComponent().append(myQualifiedName.getLastComponent())
+             : myQualifiedName;
+    }
   }
 
   public static final class QualifiedNameMatcher extends PrefixMatcher {

python/testData/completion/notImportedPackageName/CompletionForAlias/main.after.py

@@ -0,0 +1 @@
+plt.

python/testData/completion/notImportedPackageName/CompletionForAlias/main.py

@@ -0,0 +1 @@
+plt<caret>

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

@@ -0,0 +1,3 @@
+import matplotlib.pyplot as plt
+
+plt.func()

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

@@ -0,0 +1 @@
+plt.<caret>

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

@@ -0,0 +1,3 @@
+import matplotlib.pyplot as plt
+
+plt.func()

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

@@ -0,0 +1 @@
+plt.fu<caret>

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

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

python/testData/completion/notImportedQualifiedName/ImportForModuleFunction/foo/bar.py

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

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

@@ -15,6 +15,13 @@ public class PyNotImportedPackageNameCompletionTest extends PyTestCase {
     myFixture.checkResultByFile(getTestName(true) + "/main.after.py");
   }
 
+  public void testCompletionForAlias() {
+    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/notImportedPackageName/";

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

@@ -45,14 +45,28 @@ public class PyNotImportedQualifiedNameCompletionTest extends PyTestCase {
     assertFalse(matcher.prefixMatches("xxx.foo.bar.baz"));
   }
 
-  public void testImportForModuleFunctionIsAdded() {
+  public void testImportForModuleFunction() {
     myFixture.copyDirectoryToProject(getTestName(false), "");
     myFixture.configureByFile("main.py");
     myFixture.completeBasic();
     myFixture.checkResultByFile(getTestName(true) + "/main.after.py");
   }
 
-  public void testImportForModuleClassIsAdded() {
+  public void testImportForModuleClass() {
+    myFixture.copyDirectoryToProject(getTestName(false), "");
+    myFixture.configureByFile("main.py");
+    myFixture.completeBasic();
+    myFixture.checkResultByFile(getTestName(true) + "/main.after.py");
+  }
+
+  public void testImportForAlias() {
+    myFixture.copyDirectoryToProject(getTestName(false), "");
+    myFixture.configureByFile("main.py");
+    myFixture.completeBasic();
+    myFixture.checkResultByFile(getTestName(true) + "/main.after.py");
+  }
+
+  public void testImportForAliasWithNonEmptyAttribute() {
     myFixture.copyDirectoryToProject(getTestName(false), "");
     myFixture.configureByFile("main.py");
     myFixture.completeBasic();