[pycharm] PY-79552 Ensure that environment is downloaded before installing an executable

(cherry picked from commit 406e678d11aee2288bbb4a40cadb6d0744eb3dbb)

IJ-MR-159906

GitOrigin-RevId: d5505eac98a5bc71c2e9daf476e8be7fc833e991

python/python-sdk/src/com/jetbrains/python/sdk/flavors/PythonSdkFlavor.java

@@ -22,6 +22,7 @@ import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.psi.icons.PythonPsiApiIcons;
 import com.jetbrains.python.run.CommandLinePatcher;
 import com.jetbrains.python.sdk.*;
+import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -48,10 +49,11 @@ public abstract class PythonSdkFlavor<D extends PyFlavorData> {
   public static final ExtensionPointName<PythonSdkFlavor<?>> EP_NAME = ExtensionPointName.create("Pythonid.pythonSdkFlavor");
   /**
    * <code>
-   *   Python 3.11
+   * Python 3.11
    * </code>
    */
-  private static final String PYTHON_VERSION_STRING_PREFIX = "Python ";
+  @ApiStatus.Internal
+  public static final String PYTHON_VERSION_STRING_PREFIX = "Python ";
   /**
    * To prevent log pollution and slowness, we cache every {@link #isFileExecutable(String, TargetEnvironmentConfiguration)} call
    * and only log it once
@@ -65,7 +67,7 @@ public abstract class PythonSdkFlavor<D extends PyFlavorData> {
   private static final Logger LOG = Logger.getInstance(PythonSdkFlavor.class);
   /**
    * <code>
-   *   python --version
+   * python --version
    * </code>
    */
   public static final String PYTHON_VERSION_ARG = "--version";

python/src/com/jetbrains/python/sdk/PySdkToInstall.kt

@@ -52,8 +52,18 @@ fun installSdkIfNeeded(sdk: Sdk, module: Module?, existingSdks: List<Sdk>, conte
  * Generic PySdkToInstall. Compatible with all OS / CpuArch.
  */
 @Internal
-class PySdkToInstall(val installation: BinaryInstallation)
-  : ProjectJdkImpl(installation.release.title, PythonSdkType.getInstance(), "", installation.release.version) {
+class PySdkToInstall(
+  val installation: BinaryInstallation,
+) : ProjectJdkImpl(
+  installation.release.title,
+  PythonSdkType.getInstance(),
+  "",
+  /**
+   * We use [com.jetbrains.python.sdk.flavors.PythonSdkFlavor.getLanguageLevelFromVersionStringStaticSafe] to parse versions of this type
+   * of SDK. That method relies on the version string being prepended with "Python ".
+   */
+  "${PythonSdkFlavor.PYTHON_VERSION_STRING_PREFIX}${installation.release.version}"
+) {
 
   /**
    * Customize [renderer], which is typically either [com.intellij.ui.ColoredListCellRenderer] or [com.intellij.ui.ColoredTreeCellRenderer].

python/src/com/jetbrains/python/sdk/add/v2/CustomNewEnvironmentCreator.kt

@@ -7,6 +7,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.ProjectJdkTable
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.ui.validation.DialogValidationRequestor
+import com.intellij.openapi.util.io.toNioPathOrNull
 import com.intellij.platform.ide.progress.ModalTaskOwner
 import com.intellij.platform.ide.progress.runWithModalProgressBlocking
 import com.intellij.platform.ide.progress.withBackgroundProgress
@@ -26,6 +27,7 @@ import com.jetbrains.python.sdk.flavors.PythonSdkFlavor
 import com.jetbrains.python.statistics.InterpreterCreationMode
 import com.jetbrains.python.statistics.InterpreterType
 import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.launch
 import org.jetbrains.annotations.ApiStatus.Internal
 import java.nio.file.Path
 
@@ -36,7 +38,6 @@ internal abstract class CustomNewEnvironmentCreator(
 ) : PythonNewEnvironmentCreator(model) {
   internal lateinit var basePythonComboBox: PythonInterpreterComboBox
 
-
   override fun buildOptions(panel: Panel, validationRequestor: DialogValidationRequestor, errorSink: ErrorSink) {
     with(panel) {
       row(message("sdk.create.custom.base.python")) {
@@ -67,7 +68,11 @@ internal abstract class CustomNewEnvironmentCreator(
   }
 
   override fun onShown() {
-    basePythonComboBox.setItems(model.baseInterpreters)
+    model.scope.launch {
+      model.baseInterpreters.collect {
+        basePythonComboBox.setItems(model.baseInterpreters)
+      }
+    }
   }
 
   override suspend fun getOrCreateSdk(moduleOrProject: ModuleOrProject): Result<Sdk, PyError> {
@@ -111,13 +116,15 @@ internal abstract class CustomNewEnvironmentCreator(
                               InterpreterCreationMode.CUSTOM)
 
   /**
-   * Creates an installation fix for executable (poetry, pipenv).
+   * Creates an installation fix for an executable (poetry, pipenv, uv, hatch).
    *
-   * 1. Checks does a `pythonExecutable` have pip.
-   * 2. If no, checks is pip is installed globally.
-   * 3. If no, downloads and installs pip from "https://bootstrap.pypa.io/get-pip.py"
-   * 4. Runs (pythonExecutable -m) pip install `package_name` --user
-   * 5. Reruns `detectExecutable`
+   * 1. Checks if the installation of the fix requires an undownloaded env.
+   * 2. If it doesn't, downloads the env and selects it.
+   * 3. Checks if `pythonExecutable` has pip.
+   * 4. If it doesn't, checks if pip is installed globally.
+   * 5. If it isn't, downloads and installs pip from "https://bootstrap.pypa.io/get-pip.py".
+   * 6. Runs `(pythonExecutable -m) pip install <package_name> --user`.
+   * 7. Reruns `detectExecutable`.
    */
   @RequiresEdt
   protected fun createInstallFix(errorSink: ErrorSink): ActionLink {
@@ -131,15 +138,31 @@ internal abstract class CustomNewEnvironmentCreator(
   }
 
   /**
-   * Installs the necessary executable in the Python environment.
+   * Downloads the selected downloadable env (if selected), then installs the necessary executable in the Python environment.
    *
    * Initiates a blocking modal progress task to:
-   * 1. Ensure pip is installed.
-   * 2. Install the executable (specified by `name`) using either a custom installation script or via pip.
+   * 1. Ensure that the environment is downloaded (if selected).
+   * 2. Ensure that pip is installed.
+   * 3. Install the executable (specified by `name`) using either a custom installation script or via pip.
    */
   @RequiresEdt
   private fun installExecutable(errorSink: ErrorSink) {
-    val pythonExecutable = model.state.baseInterpreter.get()?.homePath ?: getPythonExecutableString()
+    val baseInterpreter = model.state.baseInterpreter.get()
+
+    val installedSdk = when (baseInterpreter) {
+      is InstallableSelectableInterpreter -> installBaseSdk(baseInterpreter.sdk, model.existingSdks)
+        ?.let {
+          val installed = model.addInstalledInterpreter(it.homePath!!.toNioPathOrNull()!!, baseInterpreter.languageLevel)
+          model.state.baseInterpreter.set(installed)
+          installed
+        }
+      is DetectedSelectableInterpreter, is ExistingSelectableInterpreter, is ManuallyAddedSelectableInterpreter, null -> null
+    }
+
+    // installedSdk is null when the selected sdk isn't downloadable
+    // model.state.baseInterpreter could be null if no SDK was selected
+    val pythonExecutable = installedSdk?.homePath ?: model.state.baseInterpreter.get()?.homePath ?: getPythonExecutableString()
+
     runWithModalProgressBlocking(ModalTaskOwner.guess(), message("sdk.create.custom.venv.install.fix.title", name, "via pip")) {
       if (installationScript != null) {
         val versionArgs: List<String> = installationVersion?.let { listOf("-v", it) } ?: emptyList()
@@ -166,7 +189,6 @@ internal abstract class CustomNewEnvironmentCreator(
    */
   private val installationScript: Path? = PythonHelpersLocator.findPathInHelpers("pycharm_package_installer.py")
 
-
   /**
    * Saves the provided path to an executable in the properties of the environment
    *

python/src/com/jetbrains/python/sdk/add/v2/models.kt

@@ -19,6 +19,7 @@ import com.intellij.python.community.services.systemPython.UICustomization
 import com.intellij.python.hatch.HatchConfiguration
 import com.intellij.python.hatch.HatchVirtualEnvironment
 import com.intellij.python.hatch.getHatchService
+import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.jetbrains.python.PyBundle.message
 import com.jetbrains.python.configuration.PyConfigurableInterpreterList
 import com.jetbrains.python.errorProcessing.ErrorSink
@@ -218,6 +219,13 @@ abstract class PythonAddInterpreterModel(params: PyInterpreterModelParams, priva
     manuallyAddedInterpreters.value += ExistingSelectableInterpreter(sdk, PySdkUtil.getLanguageLevelForSdk(sdk), sdk.isSystemWide)
   }
 
+  @RequiresEdt
+  internal fun addInstalledInterpreter(homePath: Path, languageLevel: LanguageLevel): DetectedSelectableInterpreter {
+    val installedInterpreter = DetectedSelectableInterpreter(homePath.pathString, languageLevel, true)
+    _detectedInterpreters.value += installedInterpreter
+    return installedInterpreter
+  }
+
   /**
    * Given [pathToPython] returns either cleaned path (if valid) or null and reports error to [errorSink]
    */