PY-79488 Try automatically configuring SDK for all modules in a uv workspace

GitOrigin-RevId: 24b825b85f23ac7c6338585cfc87fbe374428779

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

@@ -9,12 +9,14 @@ import com.intellij.openapi.progress.runBlockingCancellable
 import com.intellij.openapi.projectRoots.ProjectJdkTable
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.projectRoots.impl.SdkConfigurationUtil
+import com.intellij.openapi.util.registry.Registry
 import com.intellij.pycharm.community.ide.impl.PyCharmCommunityCustomizationBundle
 import com.intellij.python.pyproject.PY_PROJECT_TOML
 import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
 import com.jetbrains.python.errorProcessing.MessageError
 import com.jetbrains.python.errorProcessing.PyResult
 import com.jetbrains.python.onSuccess
+import com.jetbrains.python.projectModel.uv.UvProjectModelService
 import com.jetbrains.python.sdk.*
 import com.jetbrains.python.sdk.configuration.PyProjectSdkConfigurationExtension
 import com.jetbrains.python.sdk.uv.impl.getUvExecutable
@@ -22,6 +24,7 @@ import com.jetbrains.python.sdk.uv.setupNewUvSdkAndEnvUnderProgress
 import com.jetbrains.python.venvReader.tryResolvePath
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.withContext
+import java.nio.file.Path
 
 class PyUvSdkConfiguration : PyProjectSdkConfigurationExtension {
   companion object {
@@ -58,7 +61,15 @@ class PyUvSdkConfiguration : PyProjectSdkConfigurationExtension {
   override fun supportsHeadlessModel(): Boolean = true
 
   private suspend fun createUv(module: Module): PyResult<Sdk> {
-    val workingDir = tryResolvePath(module.basePath)
+    val venvParentDir: String?
+    if (Registry.`is`("python.project.model.uv", false)) {
+      val uvWorkspace = UvProjectModelService.findWorkspace(module)
+      venvParentDir = uvWorkspace?.root?.basePath ?: module.basePath
+    }
+    else {
+      venvParentDir = module.basePath
+    }
+    val workingDir: Path? = tryResolvePath(venvParentDir)
     if (workingDir == null) {
       return PyResult.failure(MessageError("Can't determine working dir for the module"))
     }

python/src/com/jetbrains/python/projectModel/uv/UvProjectModelService.kt

@@ -40,22 +40,36 @@ object UvProjectModelService : BaseProjectModelService<UvEntitySource, UvProject
 
   override fun getEntitySourceClass(): KClass<out EntitySource> = UvEntitySource::class
 
-  fun findSameWorkspaceMembers(module: Module): List<Module> {
+  fun findWorkspace(module: Module): UvWorkspace<Module>? {
     val wsmSnapshot = module.project.workspaceModel.currentSnapshot
     val moduleEntity = wsmSnapshot.resolve(ModuleId(module.name))!!
-    return findSameWorkspaceMembers(module.project, moduleEntity)
-      .mapNotNull { it.findModule(wsmSnapshot) }
+    val workspace = findWorkspace(module.project, moduleEntity)
+    if (workspace == null) {
+      return null
+    }
+    return UvWorkspace(
+      root = workspace.root.findModule(wsmSnapshot)!!,
+      members = workspace.members.mapNotNull { it.findModule(wsmSnapshot) }.toSet(),
+    )
   }
-  
-  fun findSameWorkspaceMembers(project: Project, module: ModuleEntity): List<ModuleEntity> {
-    val workspaceName = (module.exModuleOptions?.linkedProjectId ?: module.name).split(":")[0]
+
+  fun findWorkspace(project: Project, module: ModuleEntity): UvWorkspace<ModuleEntity>? {
+    val fullName = module.exModuleOptions?.linkedProjectId
+    if (fullName == null) return null
+    val workspaceName = fullName.split(":")[0]
     val currentSnapshot = project.workspaceModel.currentSnapshot
-    return currentSnapshot.entitiesBySource { it is UvEntitySource }
-      .filterIsInstance<ModuleEntity>()
-      .filter { 
-        val externalId = it.exModuleOptions?.linkedProjectId
-        externalId != null && (externalId == workspaceName || externalId.startsWith("$workspaceName:"))
-      }
-      .toList()
+    val rootModule = currentSnapshot.resolve(ModuleId(workspaceName))
+    return UvWorkspace(
+      root = rootModule!!,
+      members = currentSnapshot.entitiesBySource { it is UvEntitySource }
+        .filterIsInstance<ModuleEntity>()
+        .filter {
+          val externalId = it.exModuleOptions?.linkedProjectId
+          externalId != null && externalId.startsWith("$workspaceName:")
+        }
+        .toSet(),
+    )
   }
+
+  data class UvWorkspace<T>(val root: T, val members: Set<T>)
 }

python/src/com/jetbrains/python/sdk/configuration/PyProjectSdkConfiguration.kt

@@ -17,6 +17,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.project.isNotificationSilentMode
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.registry.Registry
 import com.intellij.openapi.util.use
 import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
 import com.jetbrains.python.PyBundle
@@ -24,9 +25,11 @@ import com.jetbrains.python.PySdkBundle
 import com.jetbrains.python.PythonPluginDisposable
 import com.jetbrains.python.inspections.PyInspectionExtension
 import com.jetbrains.python.inspections.requirement.RunningPackagingTasksListener
+import com.jetbrains.python.projectModel.uv.UvProjectModelService
 import com.jetbrains.python.psi.PyFile
 import com.jetbrains.python.sdk.PySdkPopupFactory
 import com.jetbrains.python.sdk.configurePythonSdk
+import com.jetbrains.python.sdk.uv.isUv
 
 object PyProjectSdkConfiguration {
 
@@ -50,7 +53,20 @@ object PyProjectSdkConfiguration {
     ProgressManager.progress("")
     LOGGER.debug("Configuring sdk with ${extension.javaClass.canonicalName} extension")
 
-    supplier()?.let { setReadyToUseSdk(module.project, module, it) }
+    val sdk = supplier()
+    if (sdk != null) {
+      // TODO Move this to PyUvSdkConfiguration, show better notification
+      if (sdk.isUv && Registry.`is`("python.project.model.uv", false)) {
+        val ws = UvProjectModelService.findWorkspace(module)
+        if (ws != null) {
+          for (wsModule in ws.members + ws.root) {
+            setReadyToUseSdk(wsModule.project, wsModule, sdk)    
+          }
+          return
+        }
+      }
+      setReadyToUseSdk(module.project, module, sdk)
+    }
   }
 
   fun setReadyToUseSdk(project: Project, module: Module, sdk: Sdk) {

python/testSrc/com/jetbrains/python/projectModel/uv/PyUvSyncIntegrationTest.kt

@@ -4,7 +4,6 @@ package com.jetbrains.python.projectModel.uv
 import com.intellij.openapi.externalSystem.testFramework.fixtures.multiProjectFixture
 import com.intellij.openapi.project.Project
 import com.intellij.platform.backend.workspace.workspaceModel
-import com.intellij.platform.testFramework.assertion.collectionAssertion.CollectionAssertions
 import com.intellij.platform.testFramework.assertion.moduleAssertion.ContentRootAssertions
 import com.intellij.platform.testFramework.assertion.moduleAssertion.DependencyAssertions
 import com.intellij.platform.testFramework.assertion.moduleAssertion.ModuleAssertions
@@ -82,7 +81,7 @@ class PyUvSyncIntegrationTest {
       name = "lib2"
       dependencies = []
     """.trimIndent())
-    
+
     testRoot.createFile("packages/pkg/pyproject.toml").writeText("""
       [project]
       name = "pkg"
@@ -94,12 +93,14 @@ class PyUvSyncIntegrationTest {
 
     val virtualFileUrlManager = project.workspaceModel.getVirtualFileUrlManager()
     ModuleAssertions.assertModules(project, "main", "lib1", "lib2", "pkg")
-    val workspaceMembers = listOf(
-      project.findModule("main"),
-      project.findModule("lib1"),
-      project.findModule("lib2"),
+    val workspace = UvProjectModelService.UvWorkspace<ModuleEntity>(
+      root = project.findModule("main")!!,
+      members = setOf(
+        project.findModule("lib1")!!,
+        project.findModule("lib2")!!,
+      )
     )
-    
+
     ModuleAssertions.assertModuleEntity(project, "main") { module ->
       ContentRootAssertions.assertContentRoots(virtualFileUrlManager, module, testRoot)
       DependencyAssertions.assertDependencies(module, DependencyAssertions.INHERITED_SDK, DependencyAssertions.MODULE_SOURCE, "lib1", "lib2")
@@ -110,9 +111,7 @@ class PyUvSyncIntegrationTest {
         Assertions.assertTrue(dependency.exported)
       }
       Assertions.assertEquals("main", module.exModuleOptions?.linkedProjectId)
-      CollectionAssertions.assertEqualsUnordered(
-        workspaceMembers,
-        UvProjectModelService.findSameWorkspaceMembers(project, module)
+      Assertions.assertEquals(workspace, UvProjectModelService.findWorkspace(project, module)
       )
     }
 
@@ -123,9 +122,7 @@ class PyUvSyncIntegrationTest {
         Assertions.assertTrue(dependency.exported)
       }
       Assertions.assertEquals("main:lib1", module.exModuleOptions?.linkedProjectId)
-      CollectionAssertions.assertEqualsUnordered(
-        workspaceMembers,
-        UvProjectModelService.findSameWorkspaceMembers(project, module)
+      Assertions.assertEquals(workspace, UvProjectModelService.findWorkspace(project, module)
       )
     }
 
@@ -133,9 +130,7 @@ class PyUvSyncIntegrationTest {
       ContentRootAssertions.assertContentRoots(virtualFileUrlManager, module, testRoot.resolve("lib/lib2"))
       DependencyAssertions.assertDependencies(module, DependencyAssertions.INHERITED_SDK, DependencyAssertions.MODULE_SOURCE)
       Assertions.assertEquals("main:lib2", module.exModuleOptions?.linkedProjectId)
-      CollectionAssertions.assertEqualsUnordered(
-        workspaceMembers,
-        UvProjectModelService.findSameWorkspaceMembers(project, module)
+      Assertions.assertEquals(workspace, UvProjectModelService.findWorkspace(project, module)
       )
     }
 
@@ -143,17 +138,17 @@ class PyUvSyncIntegrationTest {
       ContentRootAssertions.assertContentRoots(virtualFileUrlManager, module, testRoot.resolve("packages/pkg"))
       DependencyAssertions.assertDependencies(module, DependencyAssertions.INHERITED_SDK, DependencyAssertions.MODULE_SOURCE)
       Assertions.assertEquals("pkg", module.exModuleOptions?.linkedProjectId)
-      CollectionAssertions.assertEqualsUnordered(
-        listOf(module),
-        UvProjectModelService.findSameWorkspaceMembers(project, module)
+      Assertions.assertEquals(
+        UvProjectModelService.UvWorkspace(module, emptySet()),
+        UvProjectModelService.findWorkspace(project, module)
       )
     }
   }
-  
+
   private fun Project.findModule(name: String): ModuleEntity? {
     return workspaceModel.currentSnapshot.resolve(ModuleId(name))
   }
-  
+
   private suspend fun syncAllProjects(project: Project) {
     multiprojectFixture.awaitProjectConfiguration(project) {
       UvProjectModelService.syncAllProjectModelRoots(project)