From 1a5ea1b01bde1fb9d0ea8e2542e97a11c7fdbeba Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Thu, 22 May 2025 13:44:38 -0700 Subject: [PATCH] PY-79488 Try automatically configuring SDK for all modules in a uv workspace GitOrigin-RevId: 24b825b85f23ac7c6338585cfc87fbe374428779 --- .../configuration/PyUvSdkConfiguration.kt | 13 +++++- .../projectModel/uv/UvProjectModelService.kt | 40 +++++++++++++------ .../PyProjectSdkConfiguration.kt | 18 ++++++++- .../uv/PyUvSyncIntegrationTest.kt | 37 ++++++++--------- 4 files changed, 72 insertions(+), 36 deletions(-) diff --git a/python/ide/impl/src/com/intellij/pycharm/community/ide/impl/configuration/PyUvSdkConfiguration.kt b/python/ide/impl/src/com/intellij/pycharm/community/ide/impl/configuration/PyUvSdkConfiguration.kt index 80d3bc3bd40a..da0aaddf8b50 100644 --- a/python/ide/impl/src/com/intellij/pycharm/community/ide/impl/configuration/PyUvSdkConfiguration.kt +++ b/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 { - 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")) } diff --git a/python/src/com/jetbrains/python/projectModel/uv/UvProjectModelService.kt b/python/src/com/jetbrains/python/projectModel/uv/UvProjectModelService.kt index 4193440abb28..385d5f67e315 100644 --- a/python/src/com/jetbrains/python/projectModel/uv/UvProjectModelService.kt +++ b/python/src/com/jetbrains/python/projectModel/uv/UvProjectModelService.kt @@ -40,22 +40,36 @@ object UvProjectModelService : BaseProjectModelService = UvEntitySource::class - fun findSameWorkspaceMembers(module: Module): List { + fun findWorkspace(module: Module): UvWorkspace? { 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 { - val workspaceName = (module.exModuleOptions?.linkedProjectId ?: module.name).split(":")[0] + + fun findWorkspace(project: Project, module: ModuleEntity): UvWorkspace? { + 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() - .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() + .filter { + val externalId = it.exModuleOptions?.linkedProjectId + externalId != null && externalId.startsWith("$workspaceName:") + } + .toSet(), + ) } + + data class UvWorkspace(val root: T, val members: Set) } diff --git a/python/src/com/jetbrains/python/sdk/configuration/PyProjectSdkConfiguration.kt b/python/src/com/jetbrains/python/sdk/configuration/PyProjectSdkConfiguration.kt index aa6efbfa7f2d..5b61d4d64a11 100644 --- a/python/src/com/jetbrains/python/sdk/configuration/PyProjectSdkConfiguration.kt +++ b/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) { diff --git a/python/testSrc/com/jetbrains/python/projectModel/uv/PyUvSyncIntegrationTest.kt b/python/testSrc/com/jetbrains/python/projectModel/uv/PyUvSyncIntegrationTest.kt index 221840ebcc2d..52a26993a339 100644 --- a/python/testSrc/com/jetbrains/python/projectModel/uv/PyUvSyncIntegrationTest.kt +++ b/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( + 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)