[python][idea] make poetry work again (PY-78930)

* replace runBlockingCancellable with pyModalBlocking in old sdk ui (they are on edt without progress job)
* correct setup associatedModulePath for new project in idea

Merge-request: IJ-MR-160193
Merged-by: Vitaly Legchilkin <Vitaly.Legchilkin@jetbrains.com>
(cherry picked from commit e4c229f758b8dd032737290a98f7768829bf6ad5)

# Conflicts:
#	community/python/src/com/jetbrains/python/sdk/poetry/poetry.kt

GitOrigin-RevId: 11d8463d2040d1d0806ebbc2f919d671896dcd74

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

@@ -125,7 +125,7 @@ internal class PyPoetrySdkConfiguration : PyProjectSdkConfigurationExtension {
         ProjectJdkTable.getInstance().allJdks,
         file,
         PythonSdkType.getInstance(),
-        PyPoetrySdkAdditionalData(),
+        PyPoetrySdkAdditionalData(module.basePath?.let { Path.of(it) }),
         suggestedSdkName(basePath)
       )
 

python/src/com/jetbrains/python/sdk/poetry/PyPoetrySdkAdditionalData.kt

@@ -3,18 +3,17 @@ package com.jetbrains.python.sdk.poetry
 
 import com.jetbrains.python.sdk.PythonSdkAdditionalData
 import org.jdom.Element
-
-/**
- * Additional Poetry data associated with an SDK.
- *
- */
+import java.nio.file.Path
 
 /**
  *  This source code is edited by @koxudaxi Koudai Aono <koxudaxi@gmail.com>
  */
 
 class PyPoetrySdkAdditionalData : PythonSdkAdditionalData {
-  constructor() : super(PyPoetrySdkFlavor)
+  constructor(associatedModulePath: Path?) : super(PyPoetrySdkFlavor) {
+    this.associatedModulePath = associatedModulePath?.toString()
+  }
+
   constructor(data: PythonSdkAdditionalData) : super(data)
 
   override fun save(element: Element) {
@@ -33,18 +32,11 @@ class PyPoetrySdkAdditionalData : PythonSdkAdditionalData {
     fun load(element: Element): PyPoetrySdkAdditionalData? =
       when {
         element.getAttributeValue(IS_POETRY) == "true" -> {
-          PyPoetrySdkAdditionalData().apply {
+          PyPoetrySdkAdditionalData(null).apply {
             load(element)
           }
         }
         else -> null
       }
-
-    /**
-     * Creates a new instance of data with copied fields.
-     */
-    @JvmStatic
-    fun copy(data: PythonSdkAdditionalData): PyPoetrySdkAdditionalData =
-      PyPoetrySdkAdditionalData(data)
   }
 }

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

@@ -52,7 +52,7 @@ suspend fun setupPoetrySdkUnderProgress(
   installPackages: Boolean,
   poetryPath: String? = null,
 ): Result<Sdk> {
-  val projectPath = newProjectPath ?: module?.basePath ?: project?.basePath
+  val projectPath = (newProjectPath ?: module?.basePath ?: project?.basePath)?.let { Path.of(it) }
                     ?: return Result.failure(FileNotFoundException("Can't find path to project or module"))
 
   val actualProject = project ?: module?.project
@@ -64,9 +64,14 @@ suspend fun setupPoetrySdkUnderProgress(
     setUpPoetry(projectPath, python, installPackages, poetryPath)
   }.getOrElse { return Result.failure(it) }
 
-  return createSdk(pythonExecutablePath, existingSdks, projectPath, suggestedSdkName(Path.of(projectPath)), PyPoetrySdkAdditionalData()).onSuccess { sdk ->
-    module?.let { sdk.setAssociationToModule(it) }
-  }
+  val sdk = createSdk(
+    sdkHomePath = pythonExecutablePath,
+    existingSdks = existingSdks,
+    associatedProjectPath = projectPath.toString(),
+    suggestedSdkName = suggestedSdkName(projectPath),
+    sdkAdditionalData = PyPoetrySdkAdditionalData(projectPath)
+  )
+  return sdk
 }
 
 internal val Sdk.isPoetry: Boolean
@@ -82,13 +87,13 @@ internal fun allModules(project: Project?): List<Module> {
   }?.sortedBy { it.name } ?: emptyList()
 }
 
-private suspend fun setUpPoetry(projectPathString: String, python: String?, installPackages: Boolean, poetryPath: String? = null): Result<Path> {
+private suspend fun setUpPoetry(projectPath: Path, python: String?, installPackages: Boolean, poetryPath: String? = null): Result<Path> {
   val poetryExecutablePathString = when (poetryPath) {
     is String -> poetryPath
     else -> {
-      val pyProjectToml = withContext(Dispatchers.IO) { StandardFileSystems.local().findFileByPath(projectPathString)?.findChild(PY_PROJECT_TOML) }
+      val pyProjectToml = withContext(Dispatchers.IO) { StandardFileSystems.local().findFileByPath(projectPath.toString())?.findChild(PY_PROJECT_TOML) }
       val init = pyProjectToml?.let { getPyProjectTomlForPoetry(it) } == null
-      setupPoetry(Path.of(projectPathString), python, installPackages, init).getOrElse { return Result.failure(it) }
+      setupPoetry(projectPath, python, installPackages, init).getOrElse { return Result.failure(it) }
     }
   }
 

python/src/com/jetbrains/python/sdk/poetry/ui/PoetryPanelCreator.kt

@@ -2,7 +2,6 @@
 package com.jetbrains.python.sdk.poetry.ui
 
 import com.intellij.openapi.module.Module
-import com.intellij.openapi.progress.runBlockingCancellable
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.util.UserDataHolder
@@ -13,6 +12,7 @@ import com.jetbrains.python.sdk.isAssociatedWithModule
 import com.jetbrains.python.sdk.poetry.POETRY_ICON
 import com.jetbrains.python.sdk.poetry.detectPoetryEnvs
 import com.jetbrains.python.sdk.poetry.sdkHomes
+import com.jetbrains.python.ui.pyModalBlocking
 import java.util.function.Supplier
 
 fun createPoetryPanel(
@@ -30,7 +30,7 @@ fun createPoetryPanel(
   val panels = listOfNotNull(newPoetryPanel, existingPoetryPanel)
   val existingSdkPaths = sdkHomes(existingSdks)
   val defaultPanel = when {
-    runBlockingCancellable {
+    pyModalBlocking {
       detectPoetryEnvs(module, existingSdkPaths, project?.basePath ?: newProjectPath)
     }.any { it.isAssociatedWithModule(module) } -> existingPoetryPanel
     newPoetryPanel != null -> newPoetryPanel

python/src/com/jetbrains/python/sdk/poetry/ui/PyAddExistingPoetryEnvPanel.kt

@@ -2,7 +2,6 @@
 package com.jetbrains.python.sdk.poetry.ui
 
 import com.intellij.openapi.module.Module
-import com.intellij.openapi.progress.runBlockingCancellable
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.ui.ValidationInfo
@@ -15,6 +14,7 @@ import com.jetbrains.python.sdk.add.PySdkPathChoosingComboBox
 import com.jetbrains.python.sdk.add.PyAddSdkPanel
 import com.jetbrains.python.sdk.add.addInterpretersAsync
 import com.jetbrains.python.sdk.poetry.*
+import com.jetbrains.python.ui.pyModalBlocking
 import java.awt.BorderLayout
 import java.util.concurrent.ConcurrentHashMap
 import javax.swing.Icon
@@ -43,7 +43,7 @@ class PyAddExistingPoetryEnvPanel(
       val existingSdkPaths = sdkHomes(existingSdks)
 
       val moduleSdks = allModules(project).parallelStream().flatMap { module ->
-        val sdks = runBlockingCancellable {
+        val sdks = pyModalBlocking {
           detectPoetryEnvs(module, existingSdkPaths, module.basePath)
         }.filterNot { it.isAssociatedWithAnotherModule(module) }
 
@@ -51,7 +51,7 @@ class PyAddExistingPoetryEnvPanel(
         sdks.stream()
       }.toList()
 
-      val rootSdks = runBlockingCancellable {
+      val rootSdks = pyModalBlocking {
         detectPoetryEnvs(module, existingSdkPaths, project?.basePath ?: newProjectPath)
       }.filterNot { it.isAssociatedWithAnotherModule(module) }
 
@@ -70,7 +70,7 @@ class PyAddExistingPoetryEnvPanel(
     return when (val sdk = sdkComboBox.selectedSdk) {
       is PyDetectedSdk -> {
         val mappedModule = sdkToModule[sdk.name] ?: module
-        runBlockingCancellable {
+        pyModalBlocking {
           setupPoetrySdkUnderProgress(project, mappedModule, existingSdks, newProjectPath,
                                       getPythonExecutable(sdk.name), false, sdk.name).onSuccess {
             PySdkSettings.instance.preferredVirtualEnvBaseSdk = getPythonExecutable(sdk.name)

python/src/com/jetbrains/python/sdk/poetry/ui/PyAddNewPoetryPanel.kt

@@ -6,7 +6,6 @@ import com.intellij.ide.util.PropertiesComponent
 import com.intellij.openapi.components.service
 import com.intellij.openapi.fileChooser.FileChooserDescriptorFactory
 import com.intellij.openapi.module.Module
-import com.intellij.openapi.progress.runBlockingCancellable
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.openapi.ui.ComboBox
@@ -33,6 +32,7 @@ import com.jetbrains.python.sdk.basePath
 import com.jetbrains.python.sdk.poetry.*
 import com.jetbrains.python.statistics.InterpreterTarget
 import com.jetbrains.python.statistics.InterpreterType
+import com.jetbrains.python.ui.pyModalBlocking
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
@@ -135,7 +135,7 @@ class PyAddNewPoetryPanel(
 
   override fun getOrCreateSdk(): Sdk? {
     PropertiesComponent.getInstance().poetryPath = poetryPathField.text.nullize()
-    return runBlockingCancellable {
+    return pyModalBlocking {
       setupPoetrySdkUnderProgress(project, selectedModule, existingSdks, newProjectPath,
                                   baseSdkField.selectedSdk.homePath, installPackagesCheckBox.isSelected).onSuccess {
         PySdkSettings.instance.preferredVirtualEnvBaseSdk = baseSdkField.selectedSdk.homePath

python/src/com/jetbrains/python/ui/ModalUtil.kt

@@ -1,9 +1,9 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.jetbrains.python.ui
 
-import com.intellij.platform.ide.progress.runWithModalProgressBlocking
 import com.intellij.platform.ide.progress.ModalTaskOwner
 import com.intellij.platform.ide.progress.TaskCancellation
+import com.intellij.platform.ide.progress.runWithModalProgressBlocking
 import com.intellij.util.concurrency.annotations.RequiresBlockingContext
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.jetbrains.python.PySdkBundle
@@ -13,7 +13,7 @@ import com.jetbrains.python.PySdkBundle
  */
 @RequiresEdt
 @RequiresBlockingContext
-fun <T> pyModalBlocking(modalTaskOwner: ModalTaskOwner = ModalTaskOwner.guess(), code: () -> T): T =
+fun <T> pyModalBlocking(modalTaskOwner: ModalTaskOwner = ModalTaskOwner.guess(), code: suspend () -> T): T =
   runWithModalProgressBlocking(modalTaskOwner, PySdkBundle.message("python.sdk.run.wait"), TaskCancellation.nonCancellable()) {
     code.invoke()
   }