PY-79488 Keep the information about Uv workspaces in the workspace model

Piggyback on ExternalSystemModuleOptionsEntity for the time being.

GitOrigin-RevId: 899a72018e8f72ed326078ccbc1315cb59cbe227

python/src/com/jetbrains/python/projectModel/BaseProjectModelService.kt

@@ -99,22 +99,26 @@ abstract class BaseProjectModelService<E : EntitySource, P : ExternalProject> {
   ): EntityStorage {
     val fileUrlManager = project.workspaceModel.getVirtualFileUrlManager()
     val storage = MutableEntityStorage.create()
-    for (module in graph) {
+    for (extProject in graph) {
       val existingModuleEntity = project.workspaceModel.currentSnapshot
         .entitiesBySource { it == source }
         .filterIsInstance<ModuleEntity>()
-        .find { it.name == module.name }
+        .find { it.name == extProject.name }
       val existingSdkEntity = existingModuleEntity
         ?.dependencies
         ?.find { it is SdkDependency } as? SdkDependency
       val sdkDependency = existingSdkEntity ?: InheritedSdkDependency
-      storage addEntity ModuleEntity(module.name, emptyList(), source) {
+      storage addEntity ModuleEntity(extProject.name, emptyList(), source) {
         dependencies += sdkDependency
         dependencies += ModuleSourceDependency
-        for (moduleName in module.dependencies) {
+        for (moduleName in extProject.dependencies) {
           dependencies += ModuleDependency(ModuleId(moduleName.name), true, DependencyScope.COMPILE, false)
         }
-        contentRoots = listOf(ContentRootEntity(module.root.toVirtualFileUrl(fileUrlManager), emptyList(), source))
+        contentRoots = listOf(ContentRootEntity(extProject.root.toVirtualFileUrl(fileUrlManager), emptyList(), source))
+        exModuleOptions = ExternalSystemModuleOptionsEntity(source) {
+          externalSystem = systemName          
+          linkedProjectId = extProject.fullName
+        }
       }
     }
     return storage

python/src/com/jetbrains/python/projectModel/poetry/PoetryProjectModelResolver.kt

@@ -13,8 +13,12 @@ import kotlin.io.path.*
 data class PoetryProject(
   override val name: String,
   override val root: Path,
-  override val dependencies: List<ExternalProjectDependency>
-) : ExternalProject
+  override val dependencies: List<ExternalProjectDependency>,
+) : ExternalProject {
+  // Poetry projects don't have any declarative hierarchical structure
+  override val fullName: String? 
+    get() = name
+}
 
 @OptIn(ExperimentalPathApi::class)
 object PoetryProjectModelResolver : PythonProjectModelResolver<PoetryProject> {

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

@@ -22,6 +22,12 @@ interface ExternalProject {
   val name: String
   val root: Path
   val dependencies: List<ExternalProjectDependency>
+
+  /**
+   * The colon separated fully qualified name of the project in the form `root:subproject:subsubproject` 
+   * if there is the given project system supports hierarchical organization of projects (e.g., uv workspaces).
+   */
+  val fullName: String?
 }
 
 data class ExternalProjectDependency(val name: String, val path: Path)

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

@@ -19,6 +19,7 @@ data class UvProject(
   override val name: String,
   override val root: Path,
   override val dependencies: List<ExternalProjectDependency>,
+  override val fullName: String?,
   val isWorkspace: Boolean,
   val parentWorkspace: UvProject?,
 ) : ExternalProject
@@ -86,6 +87,7 @@ object UvProjectModelResolver : PythonProjectModelResolver<UvProject> {
         dependencies = it.pathDependencies.map { dep -> ExternalProjectDependency(name = dep.key, path = dep.value) },
         isWorkspace = false,
         parentWorkspace = null,
+        fullName = it.projectName,
       )
     }
 
@@ -100,6 +102,7 @@ object UvProjectModelResolver : PythonProjectModelResolver<UvProject> {
           .map { ExternalProjectDependency(name = it.key, path = it.value) },
         isWorkspace = true,
         parentWorkspace = null,
+        fullName = wsRootToml.projectName,
       )
       allUvProjects.add(wsRootProject)
       for ((_, wsMemberToml) in wsMembersByNames) {
@@ -110,6 +113,7 @@ object UvProjectModelResolver : PythonProjectModelResolver<UvProject> {
             .map { ExternalProjectDependency(name = it.key, path = it.value) },
           isWorkspace = false,
           parentWorkspace = wsRootProject,
+          fullName = "${wsRootProject.name}:${wsMemberToml.projectName}"
         ))
       }
     }

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

@@ -2,11 +2,16 @@
 package com.jetbrains.python.projectModel.uv
 
 import com.intellij.openapi.components.service
+import com.intellij.openapi.module.Module
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.NlsSafe
 import com.intellij.platform.backend.workspace.workspaceModel
+import com.intellij.platform.workspace.jps.entities.ModuleEntity
+import com.intellij.platform.workspace.jps.entities.ModuleId
+import com.intellij.platform.workspace.jps.entities.exModuleOptions
 import com.intellij.platform.workspace.storage.EntitySource
 import com.intellij.platform.workspace.storage.impl.url.toVirtualFileUrl
+import com.intellij.workspaceModel.ide.legacyBridge.findModule
 import com.jetbrains.python.projectModel.BaseProjectModelService
 import com.jetbrains.python.projectModel.ProjectModelSettings
 import com.jetbrains.python.projectModel.ProjectModelSyncListener
@@ -34,4 +39,23 @@ object UvProjectModelService : BaseProjectModelService<UvEntitySource, UvProject
   }
 
   override fun getEntitySourceClass(): KClass<out EntitySource> = UvEntitySource::class
+
+  fun findSameWorkspaceMembers(module: Module): List<Module> {
+    val wsmSnapshot = module.project.workspaceModel.currentSnapshot
+    val moduleEntity = wsmSnapshot.resolve(ModuleId(module.name))!!
+    return findSameWorkspaceMembers(module.project, moduleEntity)
+      .mapNotNull { it.findModule(wsmSnapshot) }
+  }
+  
+  fun findSameWorkspaceMembers(project: Project, module: ModuleEntity): List<ModuleEntity> {
+    val workspaceName = (module.exModuleOptions?.linkedProjectId ?: module.name).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()
+  }
 }

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

@@ -4,9 +4,13 @@ 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
+import com.intellij.platform.workspace.jps.entities.ModuleEntity
+import com.intellij.platform.workspace.jps.entities.ModuleId
+import com.intellij.platform.workspace.jps.entities.exModuleOptions
 import com.intellij.testFramework.common.timeoutRunBlocking
 import com.intellij.testFramework.junit5.RegistryKey
 import com.intellij.testFramework.junit5.TestApplication
@@ -90,6 +94,12 @@ 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"),
+    )
+    
     ModuleAssertions.assertModuleEntity(project, "main") { module ->
       ContentRootAssertions.assertContentRoots(virtualFileUrlManager, module, testRoot)
       DependencyAssertions.assertDependencies(module, DependencyAssertions.INHERITED_SDK, DependencyAssertions.MODULE_SOURCE, "lib1", "lib2")
@@ -99,6 +109,11 @@ class PyUvSyncIntegrationTest {
       DependencyAssertions.assertModuleDependency(module, "lib2") { dependency ->
         Assertions.assertTrue(dependency.exported)
       }
+      Assertions.assertEquals("main", module.exModuleOptions?.linkedProjectId)
+      CollectionAssertions.assertEqualsUnordered(
+        workspaceMembers,
+        UvProjectModelService.findSameWorkspaceMembers(project, module)
+      )
     }
 
     ModuleAssertions.assertModuleEntity(project, "lib1") { module ->
@@ -107,20 +122,39 @@ class PyUvSyncIntegrationTest {
       DependencyAssertions.assertModuleDependency(module, "pkg") { dependency ->
         Assertions.assertTrue(dependency.exported)
       }
+      Assertions.assertEquals("main:lib1", module.exModuleOptions?.linkedProjectId)
+      CollectionAssertions.assertEqualsUnordered(
+        workspaceMembers,
+        UvProjectModelService.findSameWorkspaceMembers(project, module)
+      )
     }
 
     ModuleAssertions.assertModuleEntity(project, "lib2") { module ->
       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)
+      )
     }
 
     ModuleAssertions.assertModuleEntity(project, "pkg") { module ->
       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)
+      )
     }
   }
   
-  suspend fun syncAllProjects(project: Project) {
+  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)
     }