support for uv configuration in idea sdk management

GitOrigin-RevId: 17207321c24bf3746742648efac8163b39aa88ef

python/ide/impl/src/com/intellij/pycharm/community/ide/impl/configuration/PyUvSdkConfiguration.kt

@@ -14,7 +14,8 @@ import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
 import com.jetbrains.python.sdk.*
 import com.jetbrains.python.sdk.configuration.PyProjectSdkConfigurationExtension
 import com.jetbrains.python.sdk.uv.impl.getUvExecutable
-import com.jetbrains.python.sdk.uv.setupUvSdkUnderProgress
+import com.jetbrains.python.sdk.uv.setupNewUvSdkAndEnvUnderProgress
+import com.jetbrains.python.venvReader.tryResolvePath
 
 class PyUvSdkConfiguration : PyProjectSdkConfigurationExtension {
   companion object {
@@ -51,7 +52,12 @@ class PyUvSdkConfiguration : PyProjectSdkConfigurationExtension {
   override fun supportsHeadlessModel(): Boolean = true
 
   private suspend fun createUv(module: Module): Result<Sdk> {
-    val sdk = setupUvSdkUnderProgress(ModuleOrProject.ModuleAndProject(module), ProjectJdkTable.getInstance().allJdks.toList(), null)
+    val workingDir = tryResolvePath(module.basePath)
+    if (workingDir == null) {
+      return Result.failure(IllegalStateException("Can't determine working dir for the module"))
+    }
+
+    val sdk = setupNewUvSdkAndEnvUnderProgress(module.project, workingDir, ProjectJdkTable.getInstance().allJdks.toList(), null)
     sdk.onSuccess {
       SdkConfigurationUtil.addSdk(it)
     }

python/pluginCore/resources/META-INF/plugin.xml

@@ -838,6 +838,8 @@ The Python plug-in provides smart editing for Python scripts. The feature set of
     <pyAddSdkProvider implementation="com.jetbrains.python.sdk.poetry.PyAddPoetrySdkProvider"/>
     <pythonFlavorProvider implementation="com.jetbrains.python.sdk.poetry.PyPoetrySdkFlavorProvider"/>
 
+    <!-- Uv -->
+    <pyAddSdkProvider implementation="com.jetbrains.python.sdk.uv.ui.PyAddUvSdkProvider"/>
     <pythonFlavorProvider implementation="com.jetbrains.python.sdk.uv.UvSdkFlavorProvider"/>
 
     <pythonFlavorProvider implementation="com.jetbrains.python.sdk.hatch.HatchSdkFlavorProvider"/>

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

@@ -160,7 +160,7 @@ internal abstract class CustomNewEnvironmentCreator(private val name: String, mo
    */
   internal abstract fun savePathToExecutableToProperties(path: Path?)
 
-  protected abstract suspend fun setupEnvSdk(project: Project?, module: Module?, baseSdks: List<Sdk>, projectPath: String, homePath: String?, installPackages: Boolean): Result<Sdk, PyError>
+  protected abstract suspend fun setupEnvSdk(project: Project, module: Module?, baseSdks: List<Sdk>, projectPath: String, homePath: String?, installPackages: Boolean): Result<Sdk, PyError>
 
   internal abstract suspend fun detectExecutable()
 }

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

@@ -26,7 +26,7 @@ internal class EnvironmentCreatorPip(model: PythonMutableTargetAddInterpreterMod
     PropertiesComponent.getInstance().pipEnvPath = savingPath
   }
 
-  override suspend fun setupEnvSdk(project: Project?, module: Module?, baseSdks: List<Sdk>, projectPath: String, homePath: String?, installPackages: Boolean): Result<Sdk, PyError> =
+  override suspend fun setupEnvSdk(project: Project, module: Module?, baseSdks: List<Sdk>, projectPath: String, homePath: String?, installPackages: Boolean): Result<Sdk, PyError> =
     setupPipEnvSdkUnderProgress(project, module, baseSdks, projectPath, homePath, installPackages).asPythonResult()
 
   override suspend fun detectExecutable() {

python/src/com/jetbrains/python/sdk/add/v2/hatch/HatchNewEnvironmentCreator.kt

@@ -58,8 +58,8 @@ internal class HatchNewEnvironmentCreator(
     HatchConfiguration.persistPathForTarget(hatchExecutablePath = savingPath)
   }
 
-  override suspend fun setupEnvSdk(project: Project?, module: Module?, baseSdks: List<Sdk>, projectPath: String, homePath: String?, installPackages: Boolean): Result<Sdk, PyError> {
-    project ?: return Result.failure(HatchUIError.ProjectIsNotSelected())
+  override suspend fun setupEnvSdk(project: Project, module: Module?, baseSdks: List<Sdk>, projectPath: String, homePath: String?, installPackages: Boolean): Result<Sdk, PyError> {
+    // FIXME: should work with the project only
     module ?: return Result.failure(HatchUIError.ModuleIsNotSelected())
     val selectedEnv = hatchEnvironmentProperty.get() ?: return Result.failure(HatchUIError.HatchEnvironmentIsNotSelected())
 

python/src/com/jetbrains/python/sdk/add/v2/poetry/EnvironmentCreatorPoetry.kt

@@ -66,7 +66,7 @@ internal class EnvironmentCreatorPoetry(model: PythonMutableTargetAddInterpreter
     PropertiesComponent.getInstance().poetryPath = savingPath
   }
 
-  override suspend fun setupEnvSdk(project: Project?, module: Module?, baseSdks: List<Sdk>, projectPath: String, homePath: String?, installPackages: Boolean): Result<Sdk, PyError> {
+  override suspend fun setupEnvSdk(project: Project, module: Module?, baseSdks: List<Sdk>, projectPath: String, homePath: String?, installPackages: Boolean): Result<Sdk, PyError> {
     module?.let { service<PoetryConfigService>().setInProjectEnv(it) }
     return setupPoetrySdkUnderProgress(project, module, baseSdks, projectPath, homePath, installPackages).asPythonResult()
   }

python/src/com/jetbrains/python/sdk/add/v2/uv/EnvironmentCreatorUv.kt

@@ -6,10 +6,11 @@ import com.intellij.openapi.observable.properties.ObservableMutableProperty
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.util.text.nullize
-import com.jetbrains.python.sdk.ModuleOrProject
+import com.jetbrains.python.sdk.basePath
 import com.jetbrains.python.sdk.uv.impl.setUvExecutable
-import com.jetbrains.python.sdk.uv.setupUvSdkUnderProgress
+import com.jetbrains.python.sdk.uv.setupNewUvSdkAndEnvUnderProgress
 import com.jetbrains.python.statistics.InterpreterType
+import com.jetbrains.python.venvReader.tryResolvePath
 import java.nio.file.Path
 import com.jetbrains.python.Result
 import com.jetbrains.python.errorProcessing.PyError
@@ -32,14 +33,14 @@ internal class EnvironmentCreatorUv(model: PythonMutableTargetAddInterpreterMode
     setUvExecutable(savingPath)
   }
 
-  override suspend fun setupEnvSdk(project: Project?, module: Module?, baseSdks: List<Sdk>, projectPath: String, homePath: String?, installPackages: Boolean): Result<Sdk, PyError> {
-    if (module == null) {
-      // FIXME: should not happen, proper error
-      return kotlin.Result.failure<Sdk>(Exception("module is null")).asPythonResult()
+  override suspend fun setupEnvSdk(project: Project, module: Module?, baseSdks: List<Sdk>, projectPath: String, homePath: String?, installPackages: Boolean): Result<Sdk, PyError> {
+    val workingDir = module?.basePath?.let { tryResolvePath(it) } ?: project.basePath?.let { tryResolvePath(it) }
+    if (workingDir == null) {
+      return kotlin.Result.failure<Sdk>(Exception("working dir is not specified for uv environment setup")).asPythonResult()
     }
 
     val python = homePath?.let { Path.of(it) }
-    return setupUvSdkUnderProgress(ModuleOrProject.ModuleAndProject(module), baseSdks, python).asPythonResult()
+    return setupNewUvSdkAndEnvUnderProgress(project, workingDir, baseSdks, python).asPythonResult()
   }
 
   override suspend fun detectExecutable() {

python/src/com/jetbrains/python/sdk/add/v2/uv/UvExistingEnvironmentSelector.kt

@@ -16,11 +16,13 @@ import com.jetbrains.python.sdk.add.v2.CustomExistingEnvironmentSelector
 import com.jetbrains.python.sdk.add.v2.DetectedSelectableInterpreter
 import com.jetbrains.python.sdk.add.v2.PythonMutableTargetAddInterpreterModel
 import com.jetbrains.python.sdk.associatedModulePath
+import com.jetbrains.python.sdk.basePath
 import com.jetbrains.python.sdk.isAssociatedWithModule
 import com.jetbrains.python.sdk.uv.isUv
-import com.jetbrains.python.sdk.uv.setupUvSdkUnderProgress
+import com.jetbrains.python.sdk.uv.setupExistingEnvAndSdk
 import com.jetbrains.python.statistics.InterpreterType
 import com.jetbrains.python.statistics.version
+import com.jetbrains.python.venvReader.tryResolvePath
 import java.nio.file.Path
 import kotlin.io.path.pathString
 
@@ -30,19 +32,27 @@ internal class UvExistingEnvironmentSelector(model: PythonMutableTargetAddInterp
   override val interpreterType: InterpreterType = InterpreterType.UV
 
   override suspend fun getOrCreateSdk(moduleOrProject: ModuleOrProject): Result<Sdk, PyError> {
-    val selectedInterpreterPath = selectedEnv.get()?.homePath ?: return failure("No selected interpreter")
-    val existingSdk = ProjectJdkTable.getInstance().allJdks.find { it.homePath == selectedInterpreterPath }
+    val selectedInterpreterPath = tryResolvePath(selectedEnv.get()?.homePath) ?: return failure("No selected interpreter")
+    val allSdk = ProjectJdkTable.getInstance().allJdks
+    val existingSdk = allSdk.find { it.homePath == selectedInterpreterPath.pathString }
     val associatedModule = extractModule(moduleOrProject)
+    val projectDir = tryResolvePath(associatedModule?.basePath ?: moduleOrProject.project.basePath) ?: return failure("No base path")
 
     // uv sdk in current module
     if (existingSdk != null && existingSdk.isUv && existingSdk.isAssociatedWithModule(associatedModule)) {
       return Result.success(existingSdk)
     }
 
-    val existingWorkingDir = existingSdk?.associatedModulePath?.let { Path.of(it) }
+    val existingWorkingDir = existingSdk?.associatedModulePath?.let { tryResolvePath(it) }
     val usePip = existingWorkingDir != null && !existingSdk.isUv
 
-    return setupUvSdkUnderProgress(moduleOrProject, ProjectJdkTable.getInstance().allJdks.toList(), Path.of(selectedInterpreterPath), existingWorkingDir, usePip).asPythonResult()
+    return setupExistingEnvAndSdk(
+      selectedInterpreterPath,
+      existingWorkingDir,
+      usePip,
+      projectDir,
+      allSdk.toList()
+    ).asPythonResult()
   }
 
   override suspend fun detectEnvironments(modulePath: Path) {

python/src/com/jetbrains/python/sdk/uv/UvExt.kt

@@ -2,11 +2,10 @@
 package com.jetbrains.python.sdk.uv
 
 import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.NlsSafe
-import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.openapi.vfs.toNioPathOrNull
 import com.intellij.platform.ide.progress.withBackgroundProgress
-import com.intellij.python.pyproject.PyProjectToml
+import com.intellij.python.pyproject.PY_PROJECT_TOML
 import com.intellij.util.PathUtil
 import com.jetbrains.python.PyBundle
 import com.jetbrains.python.icons.PythonIcons
@@ -15,6 +14,7 @@ import com.jetbrains.python.sdk.uv.impl.createUvCli
 import com.jetbrains.python.sdk.uv.impl.createUvLowLevel
 import java.nio.file.Path
 import javax.swing.Icon
+import kotlin.io.path.exists
 import kotlin.io.path.pathString
 
 internal val Sdk.isUv: Boolean
@@ -32,48 +32,51 @@ internal fun suggestedSdkName(basePath: Path): @NlsSafe String {
 
 val UV_ICON: Icon = PythonIcons.UV
 
-suspend fun setupUvSdkUnderProgress(
-  moduleOrProject: ModuleOrProject,
+suspend fun setupNewUvSdkAndEnvUnderProgress(
+  project: Project,
+  workingDir: Path,
   existingSdks: List<Sdk>,
-  python: Path?,
-  existingSdkWorkingDir: Path? = null,
-  usePip: Boolean = false,
+  basePython: Path?,
 ): Result<Sdk> {
+  return withBackgroundProgress(project, PyBundle.message("python.sdk.dialog.title.setting.up.uv.environment"), true) {
+    setupNewUvSdkAndEnv(workingDir, existingSdks, basePython)
+  }
+}
 
-  val (pyProjectToml, moduleWorkingDirectory) = resolveWorkingDirectory(moduleOrProject)
-  val init = pyProjectToml == null
-  val uvWorkingDir = existingSdkWorkingDir ?: moduleWorkingDirectory
-  val uv = createUvLowLevel(uvWorkingDir, createUvCli())
+suspend fun setupNewUvSdkAndEnv(
+  workingDir: Path,
+  existingSdks: List<Sdk>,
+  basePython: Path?,
+): Result<Sdk> {
+  val toml = workingDir.resolve(PY_PROJECT_TOML)
+  val init = !toml.exists()
 
-  val envExecutable =
-    if (existingSdkWorkingDir == null) {
-      withBackgroundProgress(moduleOrProject.project, PyBundle.message("python.sdk.dialog.title.setting.up.uv.environment"), true) {
-        uv.initializeEnvironment(init, python)
-      }.getOrElse {
-        return Result.failure(it)
-      }
+  val uv = createUvLowLevel(workingDir, createUvCli())
+  val envExecutable = uv.initializeEnvironment(init, basePython)
+    .getOrElse {
+      return Result.failure(it)
     }
-    else {
-      python
-    } ?: throw IllegalArgumentException("Python executable is required to setup uv environment")
 
-  val sdk = createSdk(envExecutable, existingSdks, moduleWorkingDirectory.pathString, suggestedSdkName(moduleWorkingDirectory), UvSdkAdditionalData(existingSdkWorkingDir, usePip))
+  return setupExistingEnvAndSdk(envExecutable, null, false, workingDir, existingSdks)
+}
+
+suspend fun setupExistingEnvAndSdk(
+  envExecutable: Path,
+  envWorkingDir: Path?,
+  usePip: Boolean,
+  projectDir: Path,
+  existingSdks: List<Sdk>,
+): Result<Sdk> {
+  val sdk = createSdk(
+    envExecutable,
+    existingSdks,
+    projectDir.toString(),
+    suggestedSdkName(projectDir),
+    UvSdkAdditionalData(envWorkingDir, usePip))
+
   sdk.onSuccess {
-    it.setAssociationToPath(moduleWorkingDirectory.pathString)
+    it.setAssociationToPath(projectDir.pathString)
   }
 
   return sdk
-}
-
-private suspend fun resolveWorkingDirectory(moduleOrProject: ModuleOrProject): Pair<VirtualFile?, Path> {
-  val (file, path) = when (moduleOrProject) {
-    is ModuleOrProject.ModuleAndProject -> PyProjectToml.findModuleWorkingDirectory(moduleOrProject.module)
-    is ModuleOrProject.ProjectOnly -> Pair(null, PyProjectToml.findProjectWorkingDirectory(moduleOrProject.project))
-  }
-
-  if (path == null) {
-    throw IllegalArgumentException("Path to module or working directory is required")
-  }
-
-  return file to path
 }

python/src/com/jetbrains/python/sdk/uv/ui/AddNewUvPanel.kt

@@ -26,19 +26,20 @@ import com.jetbrains.python.PyBundle
 import com.jetbrains.python.PySdkBundle
 import com.jetbrains.python.PythonModuleTypeBase
 import com.jetbrains.python.newProject.collector.InterpreterStatisticsInfo
-import com.jetbrains.python.sdk.ModuleOrProject
 import com.jetbrains.python.sdk.PySdkSettings
 import com.jetbrains.python.sdk.PythonSdkCoroutineService
-import com.jetbrains.python.sdk.add.PyAddNewEnvPanel
-import com.jetbrains.python.sdk.add.PySdkPathChoosingComboBox
-import com.jetbrains.python.sdk.add.addInterpretersAsync
+import com.jetbrains.python.sdk.add.*
 import com.jetbrains.python.sdk.basePath
-import com.jetbrains.python.sdk.uv.*
+import com.jetbrains.python.sdk.uv.UV_ICON
+import com.jetbrains.python.sdk.uv.getPyProjectTomlForUv
 import com.jetbrains.python.sdk.uv.impl.detectUvExecutable
 import com.jetbrains.python.sdk.uv.impl.getUvExecutable
 import com.jetbrains.python.sdk.uv.impl.setUvExecutable
+import com.jetbrains.python.sdk.uv.setupNewUvSdkAndEnvUnderProgress
+import com.jetbrains.python.sdk.uv.validateSdks
 import com.jetbrains.python.statistics.InterpreterTarget
 import com.jetbrains.python.statistics.InterpreterType
+import com.jetbrains.python.venvReader.tryResolvePath
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
@@ -46,6 +47,7 @@ import java.awt.BorderLayout
 import java.awt.Dimension
 import java.awt.event.ItemEvent
 import java.nio.file.Path
+import java.util.function.Supplier
 import javax.swing.Icon
 import javax.swing.JComboBox
 import javax.swing.event.DocumentEvent
@@ -77,6 +79,16 @@ class PyAddNewUvPanel(
   override val icon: Icon = UV_ICON
 
   private val moduleField: JComboBox<Module>
+
+  private val environmentLocation = TextFieldWithBrowseButton().apply {
+    addBrowseFolderListener(project, FileChooserDescriptorFactory.createSingleFolderDescriptor())
+    val path = module?.basePath ?: project?.basePath
+    text = when {
+      path != null -> path
+      else -> System.getProperty("user.home")
+    }
+  }
+
   private val baseSdkField = PySdkPathChoosingComboBox()
 
   init {
@@ -97,7 +109,7 @@ class PyAddNewUvPanel(
   }
 
   private val uvPathField = TextFieldWithBrowseButton().apply {
-    addBrowseFolderListener(null, FileChooserDescriptorFactory.createSingleFileDescriptor())
+    addBrowseFolderListener(project, FileChooserDescriptorFactory.createSingleFileDescriptor())
     val field = textField as? JBTextField ?: return@apply
     service<PythonSdkCoroutineService>().cs.launch {
       detectUvExecutable()?.let {
@@ -140,6 +152,7 @@ class PyAddNewUvPanel(
         }
         addLabeledComponent(associatedObjectLabel, moduleField)
       }
+      addLabeledComponent(PySdkBundle.message("python.venv.location.label"), environmentLocation)
       addLabeledComponent(PySdkBundle.message("python.venv.base.label"), baseSdkField)
       addComponent(installPackagesCheckBox)
       addLabeledComponent(PyBundle.message("python.sdk.uv.executable"), uvPathField)
@@ -150,11 +163,10 @@ class PyAddNewUvPanel(
   }
 
   override fun getOrCreateSdk(): Sdk? {
-    val module = selectedModule
-    val path = newProjectPath
-    val python = baseSdkField.selectedSdk.homePath
+    val path = tryResolvePath(newProjectPath ?: project?.basePath)
+    val python = tryResolvePath(baseSdkField.selectedSdk.homePath)
 
-    if (module == null || path == null || python == null) {
+    if (project == null || path == null || python == null) {
       return null
     }
 
@@ -162,8 +174,9 @@ class PyAddNewUvPanel(
     uvPath?.let {
       setUvExecutable(it)
     }
+
     val sdk = runBlockingCancellable {
-      setupUvSdkUnderProgress(ModuleOrProject.ModuleAndProject(module), existingSdks, Path.of(python))
+      setupNewUvSdkAndEnvUnderProgress(project, path, existingSdks, python)
     }
 
     sdk.onSuccess {
@@ -222,3 +235,22 @@ class PyAddNewUvPanel(
   private val projectPath: String?
     get() = newProjectPath ?: selectedModule?.basePath ?: project?.basePath
 }
+
+class PyAddUvSdkProvider : PyAddSdkProvider {
+  override fun createView(
+    project: Project?,
+    module: Module?,
+    newProjectPath: String?,
+    existingSdks: List<Sdk>,
+    context: UserDataHolder,
+  ): PyAddSdkPanel {
+    val panel = PyAddNewUvPanel(project, module, existingSdks, null, context)
+    // TODO: support for adding existing uv environments
+
+    return PyAddSdkGroupPanel(Supplier { "uv environment" }, UV_ICON, listOf(panel), panel)
+  }
+
+  private fun allowCreatingNewEnvironments(project: Project?): Boolean {
+    return project != null || !PlatformUtils.isPyCharm() || PlatformUtils.isPyCharmEducational()
+  }
+}