PY-82848 Refactor package handling to unify installed and requirement package processing in search results.

GitOrigin-RevId: c2b8ed89e4294d5607f66a2dede2f5cf3f5367fa

Author: Timur Malanin

python/src/com/jetbrains/python/packaging/packageRequirements/PythonPackageRequirementsTreeExtractor.kt

@@ -2,15 +2,14 @@
 package com.jetbrains.python.packaging.packageRequirements
 
 import com.intellij.openapi.extensions.ExtensionPointName
-import com.intellij.openapi.module.Module
 import com.intellij.openapi.projectRoots.Sdk
 import com.jetbrains.python.packaging.common.NormalizedPythonPackageName
 import com.jetbrains.python.packaging.common.PythonPackage
 import org.jetbrains.annotations.ApiStatus
 
 @ApiStatus.Internal
 interface PythonPackageRequirementsTreeExtractor {
-  suspend fun extract(pkg: PythonPackage, module: Module): PackageNode
+  suspend fun extract(pkg: PythonPackage): PackageNode
 
   companion object {
     private val treeParser = TreeParser()

python/src/com/jetbrains/python/packaging/toolwindow/PyPackagingToolWindowPanel.kt

@@ -228,7 +228,7 @@ class PyPackagingToolWindowPanel(private val project: Project) : SimpleToolWindo
     descriptionController.setPackage(selectedPackage)
   }
 
-  fun showSearchResult(installed: List<InstalledPackage>, repoData: List<PyPackagesViewData>) {
+  fun showSearchResult(installed: List<DisplayablePackage>, repoData: List<PyPackagesViewData>) {
     packageListController.showSearchResult(installed, repoData)
   }
 

python/src/com/jetbrains/python/packaging/toolwindow/PyPackagingToolWindowService.kt

@@ -10,7 +10,6 @@ import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.fileEditor.FileEditorManagerEvent
 import com.intellij.openapi.fileEditor.FileEditorManagerListener
-import com.intellij.openapi.module.Module
 import com.intellij.openapi.module.ModuleUtilCore
 import com.intellij.openapi.options.ex.SingleConfigurableEditor
 import com.intellij.openapi.project.Project
@@ -39,6 +38,7 @@ import com.jetbrains.python.packaging.toolwindow.model.*
 import com.jetbrains.python.sdk.PythonSdkUtil
 import com.jetbrains.python.sdk.pythonSdk
 import kotlinx.coroutines.*
+import org.jetbrains.annotations.ApiStatus
 import org.jetbrains.annotations.Nls
 
 @Service(Service.Level.PROJECT)
@@ -87,6 +87,68 @@ class PyPackagingToolWindowService(val project: Project, val serviceScope: Corou
     }
   }
 
+  private fun nameMatches(pkg: DisplayablePackage, query: String): Boolean {
+    val shouldUseStraightComparison = when (pkg) {
+      is InstalledPackage -> isNonPipCondaPackage(pkg.instance)
+      is RequirementPackage -> isNonPipCondaPackage(pkg.instance)
+      else -> false
+    }
+
+    return if (shouldUseStraightComparison) {
+      StringUtil.containsIgnoreCase(pkg.name, query)
+    } else {
+      StringUtil.containsIgnoreCase(normalizePackageName(pkg.name), normalizePackageName(query))
+    }
+  }
+
+  private fun isNonPipCondaPackage(pkg: PythonPackage): Boolean = pkg is CondaPackage && !pkg.installedWithPip
+
+  private fun traversePackageTree(
+    pkg: DisplayablePackage,
+    visited: MutableSet<String>,
+    matches: MutableList<RequirementPackage>,
+    query: String,
+  ) {
+    if (!visited.add(pkg.name)) return
+
+    if (pkg is RequirementPackage && nameMatches(pkg, query)) {
+      matches.add(pkg)
+    }
+
+    for (requirementPackage in pkg.getRequirements()) {
+      traversePackageTree(requirementPackage, visited, matches, query)
+    }
+  }
+
+  /**
+   * Finds all packages (both installed and requirements) that match the given query.
+   */
+  @ApiStatus.Internal
+  fun findAllMatchingPackages(query: String): List<DisplayablePackage> {
+    val matchingInstalled = installedPackages.values.filter { nameMatches(it, query) }
+    val matchingRequirements = mutableListOf<RequirementPackage>()
+    val visited = mutableSetOf<String>()
+
+    for (pkg in installedPackages.values) {
+      traversePackageTree(pkg, visited, matchingRequirements, query)
+    }
+
+    return unifyPackages(matchingInstalled, matchingRequirements)
+  }
+
+  /**
+   * Unifies packages with the same name according to the following rules:
+   * 1. If both an installed package and a requirement package have the same name, keep only the installed package.
+   * 2. If multiple requirement packages have the same name, keep only one of them.
+   */
+  private fun unifyPackages(installedPackages: List<InstalledPackage>, requirementPackages: List<RequirementPackage>): List<DisplayablePackage> {
+    return (installedPackages + requirementPackages)
+      .groupBy { it.name.lowercase() }
+      .map { (_, packages) ->
+        packages.find { it is InstalledPackage } ?: packages.first()
+      }
+  }
+
   fun handleSearch(query: String) {
     val manager = manager ?: return
     val prevSelected = toolWindowPanel?.getSelectedPackage()
@@ -95,21 +157,14 @@ class PyPackagingToolWindowService(val project: Project, val serviceScope: Corou
     if (query.isNotEmpty()) {
       searchJob?.cancel()
       searchJob = serviceScope.launch {
-        val installed = installedPackages.values.filter { pkg ->
-          when {
-            pkg.instance is CondaPackage && !pkg.instance.installedWithPip -> StringUtil.containsIgnoreCase(pkg.name, query)
-            else -> StringUtil.containsIgnoreCase(normalizePackageName(pkg.name), normalizePackageName(query))
-          }
-        }
-
-
-        val packagesFromRepos = manager.repositoryManager.searchPackages(query).map {
-          sortPackagesForRepo(it.value, query, it.key)
-        }.toList()
+        val allMatches = findAllMatchingPackages(query)
+        val packagesFromRepos = manager.repositoryManager.searchPackages(query)
+          .map { (repository, packages) -> sortPackagesForRepo(packages, query, repository) }
+          .toList()
 
         if (isActive) {
           withContext(Dispatchers.Main) {
-            toolWindowPanel?.showSearchResult(installed, packagesFromRepos + invalidRepositories)
+            toolWindowPanel?.showSearchResult(allMatches, packagesFromRepos + invalidRepositories)
             prevSelected?.name?.let { toolWindowPanel?.selectPackageName(it) }
           }
         }
@@ -147,7 +202,8 @@ class PyPackagingToolWindowService(val project: Project, val serviceScope: Corou
     handleActionCompleted(message("python.packaging.notification.deleted", selectedPackages.joinToString(", ") { it.name }))
   }
 
-  internal suspend fun initForSdk(sdk: Sdk?) {
+  @ApiStatus.Internal
+  suspend fun initForSdk(sdk: Sdk?) {
     if (sdk == null) {
       toolWindowPanel?.packageListController?.setLoadingState(false)
     }
@@ -234,7 +290,6 @@ class PyPackagingToolWindowService(val project: Project, val serviceScope: Corou
 
   suspend fun refreshInstalledPackages() {
     val sdk = currentSdk ?: return
-    val targetModule = findTargetModule(sdk) ?: return
     val manager = manager ?: return
     withContext(Dispatchers.Default) {
       val declaredPackages = manager.reloadDependencies()
@@ -245,9 +300,9 @@ class PyPackagingToolWindowService(val project: Project, val serviceScope: Corou
         processPackagesWithRequirementsTree(
           installedDeclaredPackages,
           treeExtractor,
-          targetModule
         )
-      } else {
+      }
+      else {
         emptyList()
       }
 
@@ -262,9 +317,6 @@ class PyPackagingToolWindowService(val project: Project, val serviceScope: Corou
     }
   }
 
-  private fun findTargetModule(sdk: Sdk): Module? =
-    project.modules.find { it.pythonSdk == sdk }
-
   private suspend fun findInstalledDeclaredPackages(declaredPackages: List<PythonPackage>): List<PythonPackage> =
     manager?.listInstalledPackages()?.filter {
       it.name in declaredPackages.map { pkg -> pkg.name }
@@ -273,10 +325,9 @@ class PyPackagingToolWindowService(val project: Project, val serviceScope: Corou
   private suspend fun processPackagesWithRequirementsTree(
     packages: List<PythonPackage>,
     treeExtractor: PythonPackageRequirementsTreeExtractor,
-    targetModule: Module,
   ): List<InstalledPackage> {
     return packages.mapNotNull { pkg ->
-      val tree = treeExtractor.extract(pkg, targetModule)
+      val tree = treeExtractor.extract(pkg)
       createInstalledPackageFromTree(pkg, tree)
     }
   }

python/src/com/jetbrains/python/packaging/toolwindow/PyPackagingTreeView.kt

@@ -52,7 +52,7 @@ internal class PyPackagingTreeView(
     }
   }
 
-  fun showSearchResult(installed: List<InstalledPackage>, repoData: List<PyPackagesViewData>) {
+  fun showSearchResult(installed: List<DisplayablePackage>, repoData: List<PyPackagesViewData>) {
     updatePackages(installed, repoData)
 
     installedPackages.expand()
@@ -94,7 +94,7 @@ internal class PyPackagingTreeView(
     container.scrollRectToVisible(Rectangle(0, 0))
   }
 
-  private fun updatePackages(installed: List<InstalledPackage>, repoData: List<PyPackagesViewData>) {
+  private fun updatePackages(installed: List<DisplayablePackage>, repoData: List<PyPackagesViewData>) {
     val sortedInstalled = installed.sortedBy { it.name }
     installedPackages.tree.items = sortedInstalled
     updateExistingRepository(installedPackages, sortedInstalled)

python/src/com/jetbrains/python/packaging/toolwindow/packages/PyPackagesListController.kt

@@ -50,7 +50,7 @@ internal class PyPackagesListController(val project: Project, val controller: Py
 
   override fun dispose() {}
 
-  fun showSearchResult(installed: List<InstalledPackage>, repoData: List<PyPackagesViewData>) {
+  fun showSearchResult(installed: List<DisplayablePackage>, repoData: List<PyPackagesViewData>) {
     tablesView.showSearchResult(installed, repoData)
     setLoadingState(false)
   }

python/src/com/jetbrains/python/poetry/packaging/PoetryPackageRequirementsTreeExtractor.kt

@@ -2,7 +2,6 @@
 package com.jetbrains.python.poetry.packaging
 
 import com.intellij.openapi.diagnostic.thisLogger
-import com.intellij.openapi.module.Module
 import com.intellij.openapi.projectRoots.Sdk
 import com.jetbrains.python.packaging.common.NormalizedPythonPackageName
 import com.jetbrains.python.packaging.common.PythonPackage
@@ -18,7 +17,7 @@ import com.jetbrains.python.sdk.poetry.runPoetryWithSdk
  */
 internal class PoetryPackageRequirementsTreeExtractor(private val sdk: Sdk) : PythonPackageRequirementsTreeExtractor {
 
-  override suspend fun extract(pkg: PythonPackage, module: Module): PackageNode {
+  override suspend fun extract(pkg: PythonPackage): PackageNode {
     val data = runPoetryWithSdk(sdk, "show", "--tree", pkg.name).getOr {
       thisLogger().info("extracting requirements for package ${pkg.name}: error. Output: \n${it.error}")
       return PackageNode(NormalizedPythonPackageName.from(pkg.name))

python/src/com/jetbrains/python/uv/packaging/UvPackageRequirementsTreeExtractor.kt

@@ -2,15 +2,13 @@
 package com.jetbrains.python.uv.packaging
 
 import com.intellij.openapi.diagnostic.thisLogger
-import com.intellij.openapi.module.Module
 import com.intellij.openapi.projectRoots.Sdk
 import com.jetbrains.python.packaging.common.NormalizedPythonPackageName
 import com.jetbrains.python.packaging.common.PythonPackage
 import com.jetbrains.python.packaging.packageRequirements.PackageNode
 import com.jetbrains.python.packaging.packageRequirements.PythonPackageRequirementsTreeExtractor
 import com.jetbrains.python.packaging.packageRequirements.PythonPackageRequirementsTreeExtractor.Companion.parseTree
 import com.jetbrains.python.packaging.packageRequirements.PythonPackageRequirementsTreeExtractorProvider
-import com.jetbrains.python.sdk.basePath
 import com.jetbrains.python.sdk.uv.UvSdkAdditionalData
 import com.jetbrains.python.sdk.uv.impl.createUvCli
 import com.jetbrains.python.sdk.uv.impl.createUvLowLevel
@@ -19,23 +17,26 @@ import java.nio.file.Path
 
 internal class UvPackageRequirementsTreeExtractor(private val uvWorkingDirectory: Path?) : PythonPackageRequirementsTreeExtractor {
 
-  override suspend fun extract(pkg: PythonPackage, module: Module): PackageNode {
-    val uvWorkingDirectory = uvWorkingDirectory ?: Path.of(module.basePath!!)
-    val uv = createUvLowLevel(uvWorkingDirectory, createUvCli())
-    val out = uv.listPackageRequirementsTree(pkg).getOr {
-      thisLogger().info("extracting requires for package ${pkg.name}: error. Output: \n${it.error}")
-      return PackageNode(NormalizedPythonPackageName.from(pkg.name), mutableListOf())
+  override suspend fun extract(pkg: PythonPackage): PackageNode {
+    val workingDir = uvWorkingDirectory ?: return createLeafPackageNode(pkg.name)
+    val uvInstance = createUvLowLevel(workingDir, createUvCli())
+    val requirementsOutput = uvInstance.listPackageRequirementsTree(pkg).getOr {
+      thisLogger().info("extracting requires for package $pkg.name: error. Output: \n${it.error}")
+      return createLeafPackageNode(pkg.name)
     }
 
-    return parseTree(out.lines())
+    return parseTree(requirementsOutput.lines())
   }
+
+  private fun createLeafPackageNode(packageName: String): PackageNode =
+    PackageNode(NormalizedPythonPackageName.from(packageName), mutableListOf())
 }
 
 
 private class UvPackageRequirementsTreeExtractorProvider : PythonPackageRequirementsTreeExtractorProvider {
   override fun createExtractor(sdk: Sdk): PythonPackageRequirementsTreeExtractor? {
     if (!sdk.isUv) return null
     val data = sdk.sdkAdditionalData as? UvSdkAdditionalData ?: return null
-    return UvPackageRequirementsTreeExtractor( data.uvWorkingDirectory)
+    return UvPackageRequirementsTreeExtractor(data.uvWorkingDirectory)
   }
 }