PY-54269 Set up cross-module dependencies using the existing "transferred roots" mechanism

Namely, whenever there is a `sys.path` entry or a user-added path pointing
to the root of another project module, instead of configuring a source root
for it, we now set up the corresponding module dependency and save it
as a "transferred root" in additional data alongside extra directories
configured as source roots.

This POC solution has a number of limitations:
- A module dependency is set up only if the corresponding module
has already been added to the project. Modules are added using
"Attach to project" as before.
- Since detecting modules is happening on the interpreter introspection
phase, if project A depends on project B, once the module B is actually
configured as an IJ module, one needs to restart SDK update for project
A's SDK to detect the dependency.
- Detecting module dependencies this way requires having a configured
project interpreter for each dependent module, which might not be
practical for large projects, such as grazie-ml.
- It doesn't play well with manual editing of module dependencies in
settings.

Overall, this workaround works well with an existing multi-module project,
where a new subproject with its own SDK and module dependencies needs to
be added as an IJ module. But ideally we need to rely on the actual
project configuration files, not on runtime values of `sys.path`, and
make the workflow of configuring all project modules at once simpler.

A registry key "python.detect.cross.module.dependencies", disabled by default,
activates the feature.

As part of the feature, I also disabled configuring source roots for
`sys.path` entries not belonging to the current module. It leads
to a confusing situation where multiple modules have a content root
(with a single source root) pointing to the same directory in another
module. In a multi-module setting, it can happen if several modules
depend on some subproject that has not yet been configured as an IJ
module.

GitOrigin-RevId: 47e832edeeea21f5206c981bd604a6012100e9da

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

@@ -646,6 +646,10 @@ The Python plug-in provides smart editing for Python scripts. The feature set of
                  restartRequired="false"
                  description="Disables automatic updating of PyPI package cache and ranking on project startup."/>
 
+    <!-- Setting up cross-module dependencies -->
+    <registryKey key="python.detect.cross.module.dependencies" defaultValue="false"
+                 description="Try to detect and automatically set-up module dependencies in a multi-module project"/>
+
     <feedback.idleFeedbackSurvey implementation="com.jetbrains.python.statistics.feedback.PythonJobSurvey"/>
     <statistics.counterUsagesCollector implementationClass="com.jetbrains.python.statistics.feedback.PythonJobStatisticsCollector"/>
     <backgroundPostStartupActivity implementation="com.jetbrains.python.statistics.feedback.PythonFirstLaunchChecker"/>

python/python-psi-impl/src/com/jetbrains/python/psi/PyUtil.java

@@ -1024,17 +1024,30 @@ public final class PyUtil {
       () -> {
         final ModifiableRootModel model = ModuleRootManager.getInstance(module).getModifiableModel();
         for (VirtualFile root : roots) {
-          boolean added = false;
           for (ContentEntry entry : model.getContentEntries()) {
             final VirtualFile file = entry.getFile();
             if (file != null && VfsUtilCore.isAncestor(file, root, true)) {
               entry.addSourceFolder(root.getUrl(), JavaSourceRootType.SOURCE, true);
-              added = true;
             }
           }
+        }
+        model.commit();
+      }
+    );
+  }
 
-          if (!added) {
-            model.addContentEntry(root).addSourceFolder(root.getUrl(), JavaSourceRootType.SOURCE, true);
+  @RequiresEdt
+  public static void addModuleDependencies(@NotNull Module module, @NotNull Collection<Module> dependencies) {
+    if (dependencies.isEmpty()) {
+      return;
+    }
+
+    ApplicationManager.getApplication().runWriteAction(
+      () -> {
+        ModifiableRootModel model = ModuleRootManager.getInstance(module).getModifiableModel();
+        for (Module dependency : dependencies) {
+          if (dependency != module && model.findModuleOrderEntry(dependency) == null) {
+            model.addModuleOrderEntry(dependency);
           }
         }
         model.commit();
@@ -1057,9 +1070,25 @@ public final class PyUtil {
               entry.removeSourceFolder(folder);
             }
           }
+        }
+        model.commit();
+      }
+    );
+  }
 
-          if (roots.contains(entry.getFile()) && entry.getSourceFolders().length == 0) {
-            model.removeContentEntry(entry);
+  @RequiresEdt
+  public static void removeModuleDependencies(@NotNull Module module, @NotNull Collection<Module> dependencies) {
+    if (dependencies.isEmpty()) {
+      return;
+    }
+
+    ApplicationManager.getApplication().runWriteAction(
+      () -> {
+        final ModifiableRootModel model = ModuleRootManager.getInstance(module).getModifiableModel();
+        for (Module dependency : dependencies) {
+          ModuleOrderEntry moduleOrderEntry = model.findModuleOrderEntry(dependency);
+          if (moduleOrderEntry != null) {
+            model.removeOrderEntry(moduleOrderEntry);
           }
         }
         model.commit();

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

@@ -1,14 +1,19 @@
 // Copyright 2000-2021 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.jetbrains.python.sdk
 
+import com.google.common.collect.MultimapBuilder
+import com.google.common.collect.SetMultimap
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.application.runWriteAction
+import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.module.ModuleManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.roots.ModuleRootManager
+import com.intellij.openapi.util.registry.Registry
+import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.jetbrains.python.psi.PyUtil
 
@@ -26,6 +31,16 @@ fun removeTransferredRootsFromModulesWithSdk(project: Project, sdk: Sdk) {
   updateRootsForModulesWithSdk(project, sdk, ::removeTransferredRoots)
 }
 
+private fun updateRootsForModulesWithSdk(project: Project, sdk: Sdk?, action: (Module, Sdk) -> Unit) {
+  if (sdk == null) {
+    return
+  }
+
+  for (module in runReadAction { ModuleManager.getInstance(project).modules }) {
+    action(module, sdk)
+  }
+}
+
 /**
  * Applies [transferRoots] to all modules inheriting python sdk from the [project].
  */
@@ -61,34 +76,66 @@ fun setPathsToTransfer(sdk: Sdk, roots: Set<VirtualFile>) {
 }
 
 /**
- * Turns [getPathsToTransfer] result into [module] source roots if [module] python sdk is [sdk].
+ * Turns [getPathsToTransfer] result into [module] source roots and dependencies if [module] python sdk is [sdk].
  */
 fun transferRoots(module: Module, sdk: Sdk?) {
   if (sdk != null && module.pythonSdk == sdk) {
     runInEdt {
-      PyUtil.addSourceRoots(module, getPathsToTransfer(sdk))
+      val transferredRoots = getPathsToTransfer(sdk)
+      val newTransferredRoots = TransferredRootsDetector(module.project).detect(module, transferredRoots, transferredRoots)
+      addTransferredRoots(module, newTransferredRoots)
     }
   }
 }
 
+private fun addTransferredRoots(module: Module, newTransferredRoots: ModuleTransferredRoots) {
+  ModuleRootAndDepOps.LOG.info("Adding source roots ${newTransferredRoots.sourceRoots} to module ${module}")
+  PyUtil.addSourceRoots(module, newTransferredRoots.sourceRoots)
+  ModuleRootAndDepOps.LOG.info("Adding dependencies ${newTransferredRoots.dependencies} to module ${module}")
+  PyUtil.addModuleDependencies(module, newTransferredRoots.dependencies)
+}
+
 /**
- * Removes [getPathsToTransfer] result from [module] source roots if [module] python sdk is [sdk].
+ * Removes [getPathsToTransfer] result from [module] source roots and dependencies if [module] python sdk is [sdk].
  */
 fun removeTransferredRoots(module: Module, sdk: Sdk?) {
   if (sdk != null && module.pythonSdk == sdk) {
     runInEdt {
-      PyUtil.removeSourceRoots(module, getPathsToTransfer(sdk))
+      val transferredRoots = getPathsToTransfer(sdk)
+      val newTransferredRoots = TransferredRootsDetector(module.project).detect(module, transferredRoots, transferredRoots)
+      removeTransferredRoots(module, newTransferredRoots)
     }
   }
 }
 
-private fun updateRootsForModulesWithSdk(project: Project, sdk: Sdk?, action: (Module, Sdk) -> Unit) {
-  if (sdk == null) {
-    return
-  }
+private fun removeTransferredRoots(module: Module, newTransferredRoots: ModuleTransferredRoots) {
+  ModuleRootAndDepOps.LOG.info("Removing source roots ${newTransferredRoots.sourceRoots} from module ${module}")
+  PyUtil.removeSourceRoots(module, newTransferredRoots.sourceRoots)
+  ModuleRootAndDepOps.LOG.info("Removing dependencies ${newTransferredRoots.dependencies} from module ${module}")
+  PyUtil.removeModuleDependencies(module, newTransferredRoots.dependencies)
+}
 
-  for (module in runReadAction { ModuleManager.getInstance(project).modules }) {
-    action(module, sdk)
+fun updateTransferredRoots(project: Project, sdk: Sdk, newInProjectPaths: Set<VirtualFile>) {
+  val rootsDetector = TransferredRootsDetector(project)
+  val modulesWithThisSdk = rootsDetector.projectModules.filter { it.pythonSdk == sdk }
+  val oldTransferredRoots = getPathsToTransfer(sdk)
+  val oldTransferredRootsStructure = modulesWithThisSdk.associateWith { rootsDetector.detect(it, oldTransferredRoots, oldTransferredRoots) }
+  val newTransferredRootsStructure = modulesWithThisSdk.associateWith { rootsDetector.detect(it, newInProjectPaths, oldTransferredRoots) }
+  val newTransferredRoots = newTransferredRootsStructure.values
+    .flatMap { it.sourceRoots + it.dependencies.mapNotNull { it.baseDir } }
+    .toSet()
+
+  if (oldTransferredRoots != newTransferredRoots) {
+    for ((module, newModuleRoots) in newTransferredRootsStructure.entries) {
+      runInEdt {
+        val oldModuleRoots = oldTransferredRootsStructure[module]
+        if (oldModuleRoots != null) {
+          removeTransferredRoots(module, oldModuleRoots)
+        }
+        addTransferredRoots(module, newModuleRoots)
+      }
+    }
+    setPathsToTransfer(sdk, newTransferredRoots)
   }
 }
 
@@ -106,3 +153,48 @@ private fun updateRootsForModulesWithInheritedSdk(project: Project, sdk: Sdk?, a
   }
 }
 
+private object ModuleRootAndDepOps {
+  val LOG = thisLogger()
+}
+
+private data class ModuleTransferredRoots(val sourceRoots: Set<VirtualFile>, val dependencies: Set<Module>)
+
+private class TransferredRootsDetector(private val project: Project) {
+  val projectModules: List<Module> = runReadAction { ModuleManager.getInstance(project).modules }.toList()
+  val moduleToContentRoots: SetMultimap<Module, VirtualFile> = MultimapBuilder.hashKeys().hashSetValues().build<Module, VirtualFile>()
+  val moduleToSourceRoots: SetMultimap<Module, VirtualFile> = MultimapBuilder.hashKeys().hashSetValues().build<Module, VirtualFile>()
+
+  init {
+    for (mod in projectModules) {
+      val moduleRootManager = ModuleRootManager.getInstance(mod)
+      moduleToContentRoots.putAll(mod, moduleRootManager.contentRoots.toList())
+      moduleToSourceRoots.putAll(mod, moduleRootManager.sourceRoots.toList())
+    }
+  }
+
+  fun detect(module: Module, newPaths: Set<VirtualFile>, oldPaths: Set<VirtualFile>): ModuleTransferredRoots {
+    val sourceRoots = linkedSetOf<VirtualFile>()
+    val dependencies = linkedSetOf<Module>()
+    val manuallyMarkedModuleSourceRoots = moduleToSourceRoots[module] - oldPaths
+    val moduleContentRoots = moduleToContentRoots[module]
+    for (path in newPaths) {
+      val moduleForPath = moduleToContentRoots.entries()
+        .filter { VfsUtil.isAncestor(it.value, path, false) }
+        // Select the closest content root containing the file in case of nested modules
+        .maxByOrNull { it.value.path.length }
+        ?.key
+      if (moduleForPath == null) continue
+      if (moduleForPath == module) {
+        if (path !in moduleContentRoots && path !in manuallyMarkedModuleSourceRoots) {
+          sourceRoots.add(path)
+        }
+      }
+      else if (path == moduleForPath.baseDir) {
+        if (Registry.`is`("python.detect.cross.module.dependencies")) {
+          dependencies.add(moduleForPath)
+        }
+      }
+    }
+    return ModuleTransferredRoots(sourceRoots, dependencies)
+  }
+}

python/src/com/jetbrains/python/sdk/PythonSdkUpdater.java

@@ -28,6 +28,7 @@ import com.intellij.openapi.ui.Messages;
 import com.intellij.openapi.util.Disposer;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.io.FileUtilRt;
+import com.intellij.openapi.util.registry.Registry;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.StandardFileSystems;
 import com.intellij.openapi.vfs.VfsUtilCore;
@@ -490,30 +491,39 @@ public final class PythonSdkUpdater {
     final List<VirtualFile> localSdkPaths = buildSdkPaths(sdk, sdkRoots.first, userAddedRoots.first);
     commitSdkPathsIfChanged(sdk, localSdkPaths, forceCommit);
 
-    final var pathsToTransfer = new HashSet<VirtualFile>();
-    pathsToTransfer.addAll(sdkRoots.second);
-    pathsToTransfer.addAll(userAddedRoots.second);
-    // Presumably source and content roots that were configured manually by user, not set up automatically as "transferred"
-    HashSet<VirtualFile> nonTransferredModuleRoots = new HashSet<>(moduleRoots);
-    nonTransferredModuleRoots.removeAll(PyTransferredSdkRootsKt.getPathsToTransfer(sdk));
-    pathsToTransfer.removeAll(nonTransferredModuleRoots);
-
-    /*
-    PyTransferredSdkRootsKt#transferRoots and PyTransferredSdkRootsKt#removeTransferredRoots skip sdks
-    that are not equal to module one (editable as well).
-
-    That's why roots changes were not applied but paths to transfer were successfully set.
-
-    When current method was executed for original sdk,
-    roots changes were not applied since there were no changes in paths to transfer (they were shared with editable copy).
-     */
-    if (!pathsToTransfer.equals(PyTransferredSdkRootsKt.getPathsToTransfer(sdk))) {
+    if (Registry.is("python.detect.cross.module.dependencies")) {
+      final var transferredPathCandidates = new HashSet<VirtualFile>();
+      transferredPathCandidates.addAll(sdkRoots.second);
+      transferredPathCandidates.addAll(userAddedRoots.second);
       if (project != null) {
-        PyTransferredSdkRootsKt.removeTransferredRootsFromModulesWithSdk(project, sdk);
+        PyTransferredSdkRootsKt.updateTransferredRoots(project, sdk, transferredPathCandidates);
       }
-      PyTransferredSdkRootsKt.setPathsToTransfer(sdk, pathsToTransfer);
-      if (project != null) {
-        PyTransferredSdkRootsKt.transferRootsToModulesWithSdk(project, sdk);
+    }
+    else {
+      final var pathsToTransfer = new HashSet<VirtualFile>();
+      pathsToTransfer.addAll(sdkRoots.second);
+      pathsToTransfer.addAll(userAddedRoots.second);
+      // Presumably source and content roots that were configured manually by user, not set up automatically as "transferred"
+      HashSet<VirtualFile> nonTransferredModuleRoots = new HashSet<>(moduleRoots);
+      nonTransferredModuleRoots.removeAll(PyTransferredSdkRootsKt.getPathsToTransfer(sdk));
+      pathsToTransfer.removeAll(nonTransferredModuleRoots);
+      /*
+      PyTransferredSdkRootsKt#transferRoots and PyTransferredSdkRootsKt#removeTransferredRoots skip sdks
+      that are not equal to module one (editable as well).
+  
+      That's why roots changes were not applied but paths to transfer were successfully set.
+  
+      When current method was executed for original sdk,
+      roots changes were not applied since there were no changes in paths to transfer (they were shared with editable copy).
+      */
+      if (!pathsToTransfer.equals(PyTransferredSdkRootsKt.getPathsToTransfer(sdk))) {
+        if (project != null) {
+          PyTransferredSdkRootsKt.removeTransferredRootsFromModulesWithSdk(project, sdk);
+        }
+        PyTransferredSdkRootsKt.setPathsToTransfer(sdk, pathsToTransfer);
+        if (project != null) {
+          PyTransferredSdkRootsKt.transferRootsToModulesWithSdk(project, sdk);
+        }
       }
     }
   }

python/testSrc/com/jetbrains/python/sdk/PySdkPathsTest.kt

@@ -7,23 +7,26 @@ import com.intellij.openapi.application.runWriteActionAndWait
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.projectRoots.ProjectJdkTable
 import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.roots.ModuleRootManager
 import com.intellij.openapi.roots.OrderRootType
 import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.io.FileUtil
+import com.intellij.openapi.util.registry.Registry
 import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.testFramework.*
 import com.intellij.testFramework.assertions.Assertions.assertThat
 import com.intellij.testFramework.rules.ProjectModelRule
-import com.jetbrains.python.PyNames
 import com.jetbrains.python.PythonMockSdk
 import com.jetbrains.python.PythonPluginDisposable
-import com.jetbrains.python.configuration.PyConfigurableInterpreterList
 import com.jetbrains.python.psi.LanguageLevel
 import com.jetbrains.python.psi.PyUtil
 import org.jetbrains.annotations.NotNull
-import org.junit.*
+import org.junit.Assume.assumeTrue
+import org.junit.ClassRule
+import org.junit.Rule
+import org.junit.Test
 
 class PySdkPathsTest {
 
@@ -220,43 +223,52 @@ class PySdkPathsTest {
   }
 
   @Test
-  fun sysPathEntryOutsideSdkAndModule1ButInsideModule2() {
+  fun sysPathEntryInsideAnotherModuleDoesNotConfigureSourceRootThere() {
     val (module1, moduleRoot1) = createModule("m1")
-    val (module2, moduleRoot2) = createModule("m2")
-
-    val sdkDir = createVenvStructureInModule(moduleRoot1)
+    val (_, moduleRoot2) = createModule("m2")
 
     val entryPath1 = createSubdir(moduleRoot1)
     val entryPath2 = createSubdir(moduleRoot2)
 
-    val sdk = PythonMockSdk.create().let {
-      val properSdk = PythonMockSdk.create("Mock SDK without path", sdkDir.path, it.sdkType, LanguageLevel.getLatest())
-      registerSdk(properSdk)
-      module1.pythonSdk = properSdk
-      module2.pythonSdk = properSdk
-      return@let properSdk
+    val sdk = PythonMockSdk.create(createVenvStructureInModule(moduleRoot1).path).also {
+      registerSdk(it)
+      module1.pythonSdk = it
     }
     sdk.putUserData(PythonSdkType.MOCK_SYS_PATH_KEY, listOf(sdk.homePath, entryPath1.path, entryPath2.path))
 
     mockPythonPluginDisposable()
     updateSdkPaths(sdk)
 
-    checkRoots(sdk, module1, listOf(moduleRoot1, entryPath1, entryPath2), emptyList())
-    checkRoots(sdk, module2, listOf(moduleRoot2, entryPath1, entryPath2), emptyList())
+    checkRoots(sdk, module1, listOf(moduleRoot1, entryPath1), emptyList())
+  }
 
-    val simpleSdk = PythonMockSdk.create().also {
+  @Test
+  fun sysPathEntryPointingToAnotherModuleRootConfiguresModuleDependency() {
+    assumeTrue("The registry key 'python.detect.cross.module.dependencies' is not enabled", 
+               Registry.`is`("python.detect.cross.module.dependencies"))
+    
+    val (module1, moduleRoot1) = createModule("m1")
+    val (module2, moduleRoot2) = createModule("m2")
+
+    val sdk = PythonMockSdk.create(createVenvStructureInModule(moduleRoot1).path).also {
       registerSdk(it)
-      removeTransferredRoots(module1, sdk)
       module1.pythonSdk = it
-
-      removeTransferredRoots(module2, sdk)
-      module2.pythonSdk = it
     }
+    mockPythonPluginDisposable()
+    
+    sdk.putUserData(PythonSdkType.MOCK_SYS_PATH_KEY, listOf(moduleRoot2.path))
+    updateSdkPaths(sdk)
+    checkRoots(sdk, module1, 
+               moduleRoots = listOf(moduleRoot1), 
+               sdkRoots = listOf(), 
+               moduleDependencies = listOf(module2))
 
-    updateSdkPaths(simpleSdk)
-
-    checkRoots(simpleSdk, module1, listOf(moduleRoot1), emptyList())
-    checkRoots(simpleSdk, module2, listOf(moduleRoot2), emptyList())
+    sdk.putUserData(PythonSdkType.MOCK_SYS_PATH_KEY, listOf())
+    updateSdkPaths(sdk)
+    checkRoots(sdk, module1,
+               moduleRoots = listOf(moduleRoot1),
+               sdkRoots = listOf(),
+               moduleDependencies = listOf())
   }
 
   private fun registerSdk(it: Sdk) {
@@ -310,7 +322,13 @@ class PySdkPathsTest {
     ApplicationManager.getApplication().invokeAndWait { PlatformTestUtil.dispatchAllInvocationEventsInIdeEventQueue() }
   }
 
-  private fun checkRoots(sdk: Sdk, module: Module, moduleRoots: List<VirtualFile>, sdkRoots: List<VirtualFile>) {
+  private fun checkRoots(
+    sdk: Sdk,
+    module: Module,
+    moduleRoots: List<VirtualFile>,
+    sdkRoots: List<VirtualFile>,
+    moduleDependencies: List<Module> = emptyList<Module>(),
+  ) {
     assertThat(PyUtil.getSourceRoots(module)).containsExactlyInAnyOrder(*moduleRoots.toTypedArray())
 
     val rootProvider = sdk.rootProvider
@@ -318,6 +336,9 @@ class PySdkPathsTest {
     assertThat(classes).containsAll(sdkRoots)
     assertThat(classes).doesNotContain(*moduleRoots.toTypedArray())
     assertThat(rootProvider.getFiles(OrderRootType.SOURCES)).isEmpty()
+
+    val actualModuleDependencies = ModuleRootManager.getInstance(module).getModifiableModel().moduleDependencies
+    assertThat(actualModuleDependencies).isEqualTo(moduleDependencies.toTypedArray())
   }
 
   private fun createInSdkRoot(sdk: Sdk, relativePath: String): VirtualFile {