PY-85270 [python] Merge PyModulePackageCompletionContributor into PyClassNameCompletionContributor

Remove the now unused base class PyImportableNameCompletionContributor.

GitOrigin-RevId: 1887fb5bf35bfbc7cdd82bc31e39b0ed17e3ff35

python/python-psi-impl/resources/intellij.python.psi.impl.xml

@@ -158,10 +158,6 @@
 
     <completion.contributor language="Python" order="first" id="pyClassNameCompletionContributor"
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyClassNameCompletionContributor"/>
-    <completion.contributor language="Python" order="last"
-                            implementationClass="com.jetbrains.python.codeInsight.completion.PyModulePackageCompletionContributor"/>
-    <completion.contributor language="Python" order="first" id="pyModuleNameCompletionContributor"
-                            implementationClass="com.jetbrains.python.codeInsight.completion.PyModuleNameCompletionContributor"/>
     <completion.contributor language="Python"
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyUnresolvedModuleAttributeCompletionContributor"/>
     <completion.contributor language="Python"

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

@@ -30,12 +30,14 @@ import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil
 import com.jetbrains.python.codeInsight.imports.AddImportHelper
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
 import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.resolve.PyQualifiedNameResolveContext
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder.QualifiedNameBasedScope
 import com.jetbrains.python.psi.resolve.fromFoothold
 import com.jetbrains.python.psi.resolve.resolveQualifiedName
 import com.jetbrains.python.psi.search.PySearchUtilBase
 import com.jetbrains.python.psi.stubs.PyExportedModuleAttributeIndex
+import com.jetbrains.python.psi.stubs.PyModuleNameIndex
 import com.jetbrains.python.psi.types.TypeEvalContext
 import com.jetbrains.python.pyi.PyiFileType
 import java.util.function.LongConsumer
@@ -114,6 +116,29 @@ class PyClassNameCompletionContributor : CompletionContributor(), DumbAware {
     TimeoutUtil.run<RuntimeException>(ThrowableRunnable {
       val scope = createScope(originalFile)
       val alreadySuggested: MutableSet<QualifiedName> = HashSet()
+      
+      // Suggest importable modules and packages
+      val moduleKeys = PyModuleNameIndex.getAllKeys(originalFile.project)
+      val modulesFromIndex = moduleKeys.asSequence()
+        .filter { result.prefixMatcher.isStartMatch(it) }
+        // TODO Is there lazier API here?
+        .flatMap { PyModuleNameIndex.findByShortName(it, originalFile.project, scope) }
+        .toList()
+
+      val resolveContext = fromFoothold(originalFile)
+      modulesFromIndex.asSequence()
+        // TODO Do we need this resolve here?
+        .flatMap { resolve(it, resolveContext) }
+        .filter { PyUtil.isImportable(originalFile, it) }
+        .mapNotNull { createLookupElementBuilder(originalFile, it) }
+        .map { it.withInsertHandler(when {
+          insideStringLiteralInExtendedCompletion -> InsertHandlers.stringLiteralInsertHandler
+          else -> InsertHandlers.importingInsertHandler
+        }) }
+        .map { PrioritizedLookupElement.withPriority(it, PythonCompletionWeigher.NOT_IMPORTED_MODULE_WEIGHT.toDouble()) }
+        .forEach { result.addElement(it) }
+
+      // Suggest top-level importable names
       forEachPublicNameFromIndex(scope) { elementName: String ->
         ProgressManager.checkCanceled()
         counters.scannedNames++
@@ -164,6 +189,12 @@ class PyClassNameCompletionContributor : CompletionContributor(), DumbAware {
     })
   }
 
+  private fun resolve(module: PsiFile, resolveContext: PyQualifiedNameResolveContext): Sequence<PsiFileSystemItem> {
+    val qualifiedName = QualifiedNameFinder.findCanonicalImportPath(module, null) ?: return emptySequence()
+    return resolveQualifiedName(qualifiedName, resolveContext).asSequence()
+      .filterIsInstance<PsiFileSystemItem>()
+  }
+
   private fun getInsertHandler(
     exported: PyElement,
     position: PsiElement,

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

@@ -6,8 +6,7 @@ import com.intellij.codeInsight.completion.AutoCompletionContext
 import com.intellij.codeInsight.completion.AutoCompletionDecision
 import com.intellij.codeInsight.completion.CompletionParameters
 import com.intellij.codeInsight.completion.CompletionResultSet
-import com.intellij.codeInsight.lookup.LookupElementBuilder
-import com.intellij.codeInsight.lookup.TailTypeDecorator
+import com.intellij.codeInsight.lookup.*
 import com.intellij.openapi.module.ModuleUtilCore
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.util.io.FileUtil
@@ -22,10 +21,10 @@ import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil
 import com.jetbrains.python.psi.PyClass
 import com.jetbrains.python.psi.PyFile
 import com.jetbrains.python.psi.PyFunction
+import com.jetbrains.python.psi.icons.PythonPsiApiIcons
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder
 import com.jetbrains.python.psi.types.TypeEvalContext
 import com.jetbrains.python.sdk.legacy.PythonSdkUtil
-import com.jetbrains.python.psi.icons.PythonPsiApiIcons
 import com.jetbrains.python.sdk.skeleton.PySkeletonUtil
 import one.util.streamex.StreamEx
 
@@ -101,15 +100,18 @@ fun createLookupElementBuilder(file: PsiFile, element: PsiFileSystemItem): Looku
   val name = FileUtil.getNameWithoutExtension(element.name)
   if (!PyNames.isIdentifier(name)) return null
 
-  val importPath = QualifiedNameFinder.findCanonicalImportPath(element, file)?.removeLastComponent()
-  val tailText = if (importPath != null && importPath.componentCount > 0) " ($importPath)" else null
-
   return LookupElementBuilder.create(element, name)
-    .withTailText(tailText, true)
-    .withIcon(element.getIcon(0))
+    .withExpensiveRenderer(object : LookupElementRenderer<LookupElement>() {
+      override fun renderElement(lookupElement: LookupElement, presentation: LookupElementPresentation) {
+        val importPath = QualifiedNameFinder.findCanonicalImportPath(element, file)?.removeLastComponent()
+        presentation.setItemText(lookupElement.getLookupString())
+        presentation.setIcon(element.getIcon(0))
+        if (importPath == null) return
+        presentation.typeText = importPath.toString()
+      }
+    })
 }
 
-
 private const val ELEMENT_TYPE = 10
 private const val LOCATION = 100
 private const val PRIVATE_API = 1_000

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

@@ -1,114 +0,0 @@
-// Copyright 2000-2018 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.completion
-
-import com.intellij.codeInsight.completion.*
-import com.intellij.codeInsight.lookup.LookupElement
-import com.intellij.openapi.command.WriteCommandAction
-import com.intellij.openapi.project.DumbAware
-import com.intellij.patterns.StandardPatterns
-import com.intellij.psi.PsiDocumentManager
-import com.intellij.psi.PsiNamedElement
-import com.intellij.psi.util.PsiTreeUtil
-import com.jetbrains.python.codeInsight.imports.AddImportHelper
-import com.jetbrains.python.psi.*
-import com.jetbrains.python.psi.resolve.QualifiedNameFinder
-
-/**
- * Provides basic functionality for providing completion variants that should add an import statement or be expanded into a qualified name.
- */
-abstract class PyImportableNameCompletionContributor : CompletionContributor(), DumbAware {
-
-  protected val importingInsertHandler: InsertHandler<LookupElement> = InsertHandler { context, item ->
-    addImportForLookupElement(context, item, context.tailOffset - 1)
-  }
-
-  protected val functionInsertHandler: InsertHandler<LookupElement> = object : PyFunctionInsertHandler() {
-    override fun handleInsert(context: InsertionContext, item: LookupElement) {
-      val tailOffset = context.tailOffset - 1
-      super.handleInsert(context, item)  // adds parentheses, modifies tail offset
-      context.commitDocument()
-      addImportForLookupElement(context, item, tailOffset)
-    }
-  }
-  
-  protected val genericTypeInsertHandler: InsertHandler<LookupElement> = InsertHandler<LookupElement> { context, item ->
-    val tailOffset = context.tailOffset - 1
-    PyParameterizedTypeInsertHandler.INSTANCE.handleInsert(context, item)
-    context.commitDocument()
-    addImportForLookupElement(context, item, tailOffset)
-  } 
-
-  protected val stringLiteralInsertHandler: InsertHandler<LookupElement> = InsertHandler { context, item ->
-    val element = item.psiElement
-    if (element == null) return@InsertHandler
-    if (element is PyQualifiedNameOwner) {
-      insertStringLiteralPrefix(element.qualifiedName, element.name, context)
-    }
-    else {
-      val importPath = QualifiedNameFinder.findCanonicalImportPath(element, null)
-      if (importPath != null) {
-        insertStringLiteralPrefix(importPath.toString(), importPath.lastComponent.toString(), context)
-      }
-    }
-  }
-
-  private fun insertStringLiteralPrefix(qualifiedName: String?, name: String?, context: InsertionContext) {
-    if (qualifiedName != null && name != null) {
-      val qualifiedNamePrefix = qualifiedName.substring(0, qualifiedName.length - name.length)
-      context.document.insertString(context.startOffset, qualifiedNamePrefix)
-    }
-  }
-
-  /**
-   * Checks whether completion should be performed for a given [parameters] and delegates actual work to [doFillCompletionVariants].
-   */
-  final override fun fillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet) {
-    if (!shouldDoCompletion(parameters, result)) return
-    doFillCompletionVariants(parameters, result)
-  }
-
-  /**
-   * Subclasses should override the method to provide completion variants.
-   */
-  protected abstract fun doFillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet)
-
-  private fun shouldDoCompletion(parameters: CompletionParameters, result: CompletionResultSet): Boolean {
-    if (result.prefixMatcher.prefix.isEmpty()) {
-      result.restartCompletionOnPrefixChange(StandardPatterns.string().longerThan(0))
-      return false
-    }
-
-    val element = parameters.position
-    val parent = element.parent
-    if (parent is PyReferenceExpression && parent.isQualified) {
-      return false
-    }
-    if (parent is PyStringLiteralExpression) {
-      val prefix = parent.text.substring(0, parameters.offset - parent.textRange.startOffset)
-      if (prefix.contains(".")) {
-        return false
-      }
-    }
-
-    return PsiTreeUtil.getParentOfType(element, PyImportStatementBase::class.java) == null
-  }
-
-  fun addImportForLookupElement(context: InsertionContext, item: LookupElement, tailOffset: Int) {
-    val manager = PsiDocumentManager.getInstance(context.project)
-    val document = manager.getDocument(context.file)
-    if (document != null) {
-      manager.commitDocument(document)
-    }
-    val ref = context.file.findReferenceAt(tailOffset)
-    if (ref == null || ref.resolve() === item.psiElement) {
-      // no import statement needed
-      return
-    }
-    WriteCommandAction.writeCommandAction(context.project, context.file).run<RuntimeException> {
-      val psiElement = item.psiElement
-      if (psiElement is PsiNamedElement) {
-        AddImportHelper.addImport(psiElement, context.file, ref.element as PyElement)
-      }
-    }
-  }
-}

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

@@ -1,56 +0,0 @@
-// Copyright 2000-2018 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.completion
-
-import com.intellij.codeInsight.completion.CompletionParameters
-import com.intellij.codeInsight.completion.CompletionResultSet
-import com.intellij.codeInsight.completion.PrioritizedLookupElement
-import com.intellij.psi.PsiFile
-import com.intellij.psi.PsiFileSystemItem
-import com.jetbrains.python.psi.PyStringLiteralExpression
-import com.jetbrains.python.psi.PyUtil
-import com.jetbrains.python.psi.resolve.PyQualifiedNameResolveContext
-import com.jetbrains.python.psi.resolve.QualifiedNameFinder
-import com.jetbrains.python.psi.resolve.fromFoothold
-import com.jetbrains.python.psi.resolve.resolveQualifiedName
-import com.jetbrains.python.psi.search.PySearchUtilBase
-import com.jetbrains.python.psi.stubs.PyModuleNameIndex
-
-/**
- * Add completion variants for modules and packages.
- *
- * The completion contributor ensures that completion variants are resolvable with project source root configuration.
- * The list of completion variants does not include namespace packages (but includes their modules where appropriate).
- */
-class PyModulePackageCompletionContributor : PyImportableNameCompletionContributor() {
-
-  override fun doFillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet) {
-    if (!parameters.isExtendedCompletion) {
-      return
-    }
-
-    val targetFile = parameters.originalFile
-    val inStringLiteral = parameters.position.parent is PyStringLiteralExpression
-    val moduleKeys = PyModuleNameIndex.getAllKeys(targetFile.project)
-    val scope = PySearchUtilBase.defaultSuggestionScope(targetFile)
-    val modulesFromIndex = moduleKeys.asSequence()
-      .filter { result.prefixMatcher.prefixMatches(it) }
-      .flatMap { PyModuleNameIndex.findByShortName(it, targetFile.project, scope).asSequence() }
-      .toList()
-
-    val resolveContext = fromFoothold(targetFile)
-    modulesFromIndex.asSequence()
-      .flatMap { resolve(it, resolveContext) }
-      .filter { PyUtil.isImportable(targetFile, it) }
-      .mapNotNull { createLookupElementBuilder(targetFile, it) }
-      .map { it.withInsertHandler(if (inStringLiteral) stringLiteralInsertHandler else importingInsertHandler) }
-      .map { PrioritizedLookupElement.withPriority(it, PythonCompletionWeigher.NOT_IMPORTED_MODULE_WEIGHT.toDouble()) }
-      .forEach { result.addElement(it) }
-  }
-
-  private fun resolve(module: PsiFile, resolveContext: PyQualifiedNameResolveContext): Sequence<PsiFileSystemItem> {
-    val qualifiedName = QualifiedNameFinder.findCanonicalImportPath(module, null) ?: return emptySequence()
-    return resolveQualifiedName(qualifiedName, resolveContext).asSequence()
-      .filterIsInstance<PsiFileSystemItem>()
-  }
-
-}