PY-82800 Requirements: Improve completion for requirements

Signed-off-by: Nikita.Ashihmin <nikita.ashihmin@jetbrains.com>

GitOrigin-RevId: e6a5e67ec3c6de367fd3ee6d7cf49018655c2286

python/src/com/jetbrains/python/requirements/RequirementsCompletionUtil.kt

@@ -10,9 +10,8 @@ import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.progress.runBlockingCancellable
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
-import com.jetbrains.python.getOrThrow
 import com.jetbrains.python.packaging.management.PythonPackageManager
-import com.jetbrains.python.packaging.management.findPackageSpecification
+import com.jetbrains.python.packaging.normalizePackageName
 import com.jetbrains.python.psi.icons.PythonPsiApiIcons
 
 fun completePackageNames(project: Project, sdk: Sdk, result: CompletionResultSet) {
@@ -32,9 +31,8 @@ fun completeVersions(name: String, project: Project, sdk: Sdk, result: Completio
 
   val versions = ApplicationUtil.runWithCheckCanceled({
                                                         runBlockingCancellable {
-                                                          val packageSpecification = packageManager.findPackageSpecification(name)
-                                                                                     ?: return@runBlockingCancellable emptyList()
-                                                          repositoryManager.getPackageDetails(packageSpecification).getOrThrow().availableVersions
+                                                          repositoryManager.getVersions(normalizePackageName(name), null)
+                                                          ?: emptyList()
                                                         }
                                                       }, EmptyProgressIndicator.notNullize(ProgressManager.getInstance().progressIndicator))
 

python/src/com/jetbrains/python/requirements/RequirementsPackageNameCompletionContributor.kt

@@ -4,14 +4,20 @@ package com.jetbrains.python.requirements
 import com.intellij.codeInsight.completion.CompletionContributor
 import com.intellij.codeInsight.completion.CompletionParameters
 import com.intellij.codeInsight.completion.CompletionResultSet
-import com.jetbrains.python.requirements.psi.SimpleName
+import com.intellij.psi.impl.source.tree.LeafPsiElement
+import com.jetbrains.python.requirements.psi.RequirementsTypes
 
 class RequirementsPackageNameCompletionContributor : CompletionContributor() {
 
   override fun fillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet) {
-    if (parameters.position.parent !is SimpleName) return
+    if (parameters.originalFile !is RequirementsFile)
+      return
+    val nameElement = parameters.position
+    if ((nameElement as? LeafPsiElement)?.elementType != RequirementsTypes.IDENTIFIER)
+      return
+
     val sdk = getPythonSdk(parameters.originalFile) ?: return
 
-    completePackageNames(parameters.position.project, sdk, result)
+    completePackageNames(nameElement.project, sdk, result)
   }
 }

python/src/com/jetbrains/python/requirements/RequirementsVersionCompletionContributor.kt

@@ -4,20 +4,41 @@ package com.jetbrains.python.requirements
 import com.intellij.codeInsight.completion.CompletionContributor
 import com.intellij.codeInsight.completion.CompletionParameters
 import com.intellij.codeInsight.completion.CompletionResultSet
+import com.intellij.psi.PsiElement
+import com.intellij.psi.impl.source.tree.LeafPsiElement
 import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.psi.util.elementType
+import com.intellij.psi.util.prevLeafs
 import com.jetbrains.python.requirements.psi.NameReq
+import com.jetbrains.python.requirements.psi.RequirementsTypes
+import com.jetbrains.python.requirements.psi.VersionStmt
 
 class RequirementsVersionCompletionContributor : CompletionContributor() {
-
   override fun fillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet) {
+    if (parameters.originalFile !is RequirementsFile)
+      return
+
     val position = parameters.position
     val parent = position.parent
     val project = parameters.editor.project ?: return
 
-    if (parent is com.jetbrains.python.requirements.psi.VersionStmt) {
-      val name = PsiTreeUtil.getParentOfType(parent, NameReq::class.java)?.name?.text ?: return
-      val sdk = getPythonSdk(parameters.originalFile) ?: return
-      completeVersions(name, project, sdk, result, false)
-    }
+    val packageName = getNameBySiblings(position) ?: getNameByParent(parent) ?: return
+    val sdk = getPythonSdk(parameters.originalFile) ?: return
+    completeVersions(packageName, project, sdk, result, false)
+  }
+
+  private fun getNameByParent(parent: PsiElement?): String? {
+    if (parent !is VersionStmt)
+      return null
+    return PsiTreeUtil.getParentOfType(parent, NameReq::class.java)?.name?.text
+  }
+
+  private fun getNameBySiblings(position: PsiElement): String? {
+    if ((position as? LeafPsiElement)?.elementType != RequirementsTypes.VERSION)
+      return null
+
+    val identifier = position.prevLeafs.firstOrNull { it.elementType == RequirementsTypes.IDENTIFIER }
+
+    return identifier?.text
   }
 }