PY-72070 Introduce a CompositePythonPackageManager that combines functionality from several package managers. This change also refactors the CondaPackageManager to get rid of type check of package specifications.

(cherry picked from commit 2b41be8d768d820769a749339fac64d03925b575)

GitOrigin-RevId: 6ac187b64eb2e516e228dcfda492f2f3d7f73f49

python/src/com/jetbrains/python/packaging/conda/CompositePythonPackageManager.kt

@@ -0,0 +1,97 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.packaging.conda
+
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.projectRoots.Sdk
+import com.jetbrains.python.packaging.common.PythonPackage
+import com.jetbrains.python.packaging.common.PythonPackageSpecification
+import com.jetbrains.python.packaging.management.PythonPackageManager
+import com.jetbrains.python.packaging.management.PythonRepositoryManager
+
+class CompositePythonPackageManager(
+  project: Project,
+  sdk: Sdk,
+  private val managers: List<PythonPackageManager>,
+) : PythonPackageManager(project, sdk) {
+  @Volatile
+  override var installedPackages: List<PythonPackage> = emptyList()
+  override var repositoryManager: PythonRepositoryManager = managers.first().repositoryManager
+
+  private fun isInRepository(repositoryManager: PythonRepositoryManager, pkgName: String) =
+    repositoryManager.allPackages().contains(pkgName)
+
+  override suspend fun installPackageCommand(specification: PythonPackageSpecification, options: List<String>): Result<String> {
+    val exceptionList = mutableListOf<Throwable>()
+    for (manager in managers) {
+      repositoryManager = manager.repositoryManager
+      installedPackages = manager.installedPackages
+
+      if (!isInRepository(repositoryManager, specification.name)) continue
+      val executionResult = manager.installPackage(specification, options)
+      val executionOutcome = executionResult.getOrElse { exceptionList.add(it) }
+
+      if (executionResult.isSuccess) {
+        return Result.success(executionOutcome.toString())
+      }
+    }
+    return Result.failure(exceptionList.last())
+  }
+
+
+  override suspend fun updatePackageCommand(specification: PythonPackageSpecification): Result<String> {
+    val exceptionList = mutableListOf<Throwable>()
+
+    for (manager in managers) {
+      repositoryManager = manager.repositoryManager
+      installedPackages = manager.installedPackages
+
+      if (!isInRepository(repositoryManager, specification.name)) continue
+      val executionResult = manager.updatePackage(specification)
+      val executionOutcome = executionResult.getOrElse { exceptionList.add(it) }
+
+      if (executionResult.isSuccess) {
+        return Result.success(executionOutcome.toString())
+      }
+    }
+
+    return Result.failure(exceptionList.last())
+  }
+
+  override suspend fun uninstallPackageCommand(pkg: PythonPackage): Result<String> {
+    val exceptionList = mutableListOf<Throwable>()
+
+    for (manager in managers) {
+      repositoryManager = manager.repositoryManager
+      installedPackages = manager.installedPackages
+
+      if (!isInRepository(repositoryManager, pkg.name)) continue
+
+      val executionResult = manager.uninstallPackage(pkg)
+      val executionOutcome = executionResult.getOrElse { exceptionList.add(it) }
+      if (executionResult.isSuccess) {
+        return Result.success(executionOutcome.toString())
+      }
+    }
+
+    return Result.failure(exceptionList.last())
+  }
+
+  override suspend fun reloadPackagesCommand(): Result<List<PythonPackage>> {
+    val exceptionList = mutableListOf<Throwable>()
+
+    for (manager in managers) {
+      repositoryManager = manager.repositoryManager
+      installedPackages = manager.installedPackages
+
+      val executionResult = manager.reloadPackages()
+      val executionOutcome = executionResult.getOrElse {
+        exceptionList.add(it)
+        emptyList()
+      }
+      if (executionResult.isSuccess) {
+        return Result.success(executionOutcome)
+      }
+    }
+    return Result.failure(exceptionList.last())
+  }
+}

python/src/com/jetbrains/python/packaging/conda/CondaPackageManager.kt

@@ -12,12 +12,12 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.util.text.StringUtil
 import com.intellij.platform.ide.progress.withBackgroundProgress
-import com.intellij.platform.util.progress.withRawProgressReporter
+import com.intellij.platform.util.progress.reportRawProgress
 import com.jetbrains.python.PyBundle.message
 import com.jetbrains.python.packaging.PyExecutionException
 import com.jetbrains.python.packaging.common.PythonPackage
 import com.jetbrains.python.packaging.common.PythonPackageSpecification
-import com.jetbrains.python.packaging.pip.PipPythonPackageManager
+import com.jetbrains.python.packaging.management.PythonPackageManager
 import com.jetbrains.python.sdk.flavors.conda.PyCondaFlavorData
 import com.jetbrains.python.sdk.getOrCreateAdditionalData
 import com.jetbrains.python.sdk.targetEnvConfiguration
@@ -27,47 +27,33 @@ import org.jetbrains.annotations.ApiStatus
 import org.jetbrains.annotations.Nls
 
 @ApiStatus.Experimental
-class CondaPackageManager(project: Project, sdk: Sdk) : PipPythonPackageManager(project, sdk) {
+class CondaPackageManager(project: Project, sdk: Sdk) : PythonPackageManager(project, sdk) {
+  @Volatile
+  override var installedPackages: List<PythonPackage> = emptyList()
   override val repositoryManager = CondaRepositoryManger(project, sdk)
 
   override suspend fun installPackageCommand(specification: PythonPackageSpecification, options: List<String>): Result<String> =
-    if (specification is CondaPackageSpecification) {
-      try {
-        Result.success(runConda("install", specification.buildInstallationString() + "-y" + options, message("conda.packaging.install.progress", specification.name)))
-      }
-      catch (ex: ExecutionException) {
-        Result.failure(ex)
-      }
+    try {
+      Result.success(runConda("install", specification.buildInstallationString() + "-y" + options, message("conda.packaging.install.progress", specification.name)))
     }
-    else {
-      super.installPackageCommand(specification, options)
+    catch (ex: ExecutionException) {
+      Result.failure(ex)
     }
 
   override suspend fun updatePackageCommand(specification: PythonPackageSpecification): Result<String> =
-    if (specification is CondaPackageSpecification) {
-      try {
-        Result.success(runConda("update", listOf(specification.name, "-y"), message("conda.packaging.update.progress", specification.name)))
-      }
-      catch (ex: ExecutionException) {
-        Result.failure(ex)
-      }
+    try {
+      Result.success(runConda("update", listOf(specification.name, "-y"), message("conda.packaging.update.progress", specification.name)))
     }
-    else {
-      super.updatePackageCommand(specification)
+    catch (ex: ExecutionException) {
+      Result.failure(ex)
     }
 
-
   override suspend fun uninstallPackageCommand(pkg: PythonPackage): Result<String> =
-    if (pkg is CondaPackage && !pkg.installedWithPip) {
-      try {
-        Result.success(runConda("uninstall", listOf(pkg.name, "-y"), message("conda.packaging.uninstall.progress", pkg.name)))
-      }
-      catch (ex: ExecutionException) {
-        Result.failure(ex)
-      }
+    try {
+      Result.success(runConda("uninstall", listOf(pkg.name, "-y"), message("conda.packaging.uninstall.progress", pkg.name)))
     }
-    else {
-      super.uninstallPackageCommand(pkg)
+    catch (ex: ExecutionException) {
+      Result.failure(ex)
     }
 
   override suspend fun reloadPackagesCommand(): Result<List<PythonPackage>> =
@@ -111,7 +97,7 @@ class CondaPackageManager(project: Project, sdk: Sdk) : PipPythonPackageManager(
       val handler = CapturingProcessHandler(process, targetedCommandLine.charset, commandLineString)
 
       val result = withBackgroundProgress(project, text, true) {
-        withRawProgressReporter<ProcessOutput?> {
+        reportRawProgress<ProcessOutput?> {
           handler.runProcess(10 * 60 * 1000)
         } as ProcessOutput
       }

python/src/com/jetbrains/python/packaging/conda/CondaPackageManagerProvider.kt

@@ -3,8 +3,10 @@ package com.jetbrains.python.packaging.conda
 
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.util.registry.Registry
 import com.jetbrains.python.packaging.management.PythonPackageManager
 import com.jetbrains.python.packaging.management.PythonPackageManagerProvider
+import com.jetbrains.python.packaging.pip.PipPythonPackageManager
 import com.jetbrains.python.sdk.PythonSdkAdditionalData
 import com.jetbrains.python.sdk.flavors.conda.PyCondaFlavorData
 import org.jetbrains.annotations.ApiStatus
@@ -14,7 +16,14 @@ class CondaPackageManagerProvider : PythonPackageManagerProvider {
   override fun createPackageManagerForSdk(project: Project, sdk: Sdk): PythonPackageManager? {
     val additionalData = sdk.sdkAdditionalData as PythonSdkAdditionalData
 
-    return if (additionalData.flavorAndData.data is PyCondaFlavorData) CondaPackageManager(project, sdk)
+    val manager = if (Registry.`is`("python.packaging.conda.chain.installation")) {
+      CompositePythonPackageManager(project, sdk, listOf(CondaPackageManager(project, sdk), PipPythonPackageManager(project, sdk)))
+    }
+    else {
+      CondaPackageManager(project, sdk)
+    }
+
+    return if (additionalData.flavorAndData.data is PyCondaFlavorData) manager
     else null
   }
 }