add all modules' unassociated venvs to existing venvs dropdown

koxudaxi · Apr 29, 2021 · b702342 · b702342
1 parent 8aea363
commit b702342
Showing 1 changed file with 17 additions and 5 deletions.
diff --git a/src/com/koxudaxi/poetry/PyAddExistingPoetryEnvPanel.kt b/src/com/koxudaxi/poetry/PyAddExistingPoetryEnvPanel.kt
@@ -16,22 +16,22 @@
 package com.koxudaxi.poetry
 
 import com.intellij.openapi.module.Module
-import com.intellij.openapi.module.ModuleUtil
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.ui.ValidationInfo
 import com.intellij.openapi.util.UserDataHolder
 import com.intellij.util.ui.FormBuilder
 import com.jetbrains.python.PyBundle
 import com.jetbrains.python.PySdkBundle
-import com.jetbrains.python.PythonModuleTypeBase
 import com.jetbrains.python.sdk.*
 import com.jetbrains.python.sdk.add.PyAddSdkPanel
 import com.jetbrains.python.sdk.add.PyAddSdkView
 import com.jetbrains.python.sdk.add.PySdkPathChoosingComboBox
 import com.jetbrains.python.sdk.add.addInterpretersAsync
 import java.awt.BorderLayout
+import java.util.concurrent.ConcurrentHashMap
 import javax.swing.Icon
+import kotlin.streams.toList
 
 /**
  * @author vlan
@@ -46,6 +46,7 @@ class PyAddExistingPoetryEnvPanel(private val project: Project?,
                                   private val existingSdks: List<Sdk>,
                                   override var newProjectPath: String?,
                                   context: UserDataHolder) : PyAddSdkPanel() {
+    private var sdkToModule = ConcurrentHashMap<String, Module>()
     override val panelName: String get() = PyBundle.message("python.add.sdk.panel.name.existing.environment")
     override val icon: Icon = POETRY_ICON
     private val sdkComboBox = PySdkPathChoosingComboBox()
@@ -58,20 +59,31 @@ class PyAddExistingPoetryEnvPanel(private val project: Project?,
         add(formPanel, BorderLayout.NORTH)
         addInterpretersAsync(sdkComboBox) {
             val existingSdkPaths = sdkHomes(existingSdks)
-            detectPoetryEnvs(module, existingSdkPaths, project?.basePath ?: newProjectPath)
+            val moduleSdks = allModules(project).parallelStream().flatMap { module ->
+                val sdks = detectPoetryEnvs(module, existingSdkPaths, module.basePath)
+                        .filterNot { it.isAssociatedWithAnotherModule(module) }
+                sdks.forEach { sdkToModule.putIfAbsent(it.name, module) }
+                sdks.stream()
+            }.toList()
+            val rootSdks = detectPoetryEnvs(module, existingSdkPaths, project?.basePath ?: newProjectPath)
                     .filterNot { it.isAssociatedWithAnotherModule(module) }
+            val moduleSdkPaths = moduleSdks.map { it.name }.toSet()
+            val sdks = rootSdks.filterNot { moduleSdkPaths.contains(it.name) } + moduleSdks
+            sdks.sortedBy { it.name }
         }
     }
 
     override fun validateAll(): List<ValidationInfo> = listOfNotNull(validateSdkComboBox(sdkComboBox, this))
 
     override fun getOrCreateSdk(): Sdk? {
         return when (val sdk = sdkComboBox.selectedSdk) {
-            is PyDetectedSdk ->
-                setupPoetrySdkUnderProgress(project, module, existingSdks, newProjectPath,
+            is PyDetectedSdk -> {
+                val mappedModule = sdkToModule[sdk.name] ?: module
+                setupPoetrySdkUnderProgress(project, mappedModule, existingSdks, newProjectPath,
                         getPythonExecutable(sdk.name), false, sdk.name)?.apply {
                     PySdkSettings.instance.preferredVirtualEnvBaseSdk = getPythonExecutable(sdk.name)
                 }
+            }
             else -> sdk
         }
     }