PY-86857: Show module names in the execution console during SDK configuration.

SDK creator (`com/jetbrains/python/sdk/configuration/CreateSdkInfo.kt:19`) creates SDK, and we set `TraceContext` so every single call it does to `ExecutionService` (to run external tool for example) is accomplished with a module name.


Merge-request: IJ-MR-187827
Merged-by: Ilya Kazakevich <ilya.kazakevich@jetbrains.com>

GitOrigin-RevId: 86823980c1d7b0ee7901545617d6f904c55ecb7f

python/ide/impl/src/com/intellij/pycharm/community/ide/impl/PythonSdkConfigurator.kt

@@ -42,6 +42,7 @@ import com.jetbrains.python.sdk.configuration.PyProjectSdkConfiguration.setReady
 import com.jetbrains.python.sdk.configuration.PyProjectSdkConfiguration.setSdkUsingCreateSdkInfo
 import com.jetbrains.python.sdk.configuration.PyProjectSdkConfiguration.suppressTipAndInspectionsFor
 import com.jetbrains.python.sdk.configuration.PyProjectSdkConfigurationExtension
+import com.jetbrains.python.sdk.configuration.getSdkCreator
 import com.jetbrains.python.sdk.impl.PySdkBundle
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -241,11 +242,11 @@ class PythonSdkConfigurator : DirectoryProjectConfigurator {
     if (fallback == null) {
       return false
     }
-    val sdkCreator = fallback.checkEnvironmentAndPrepareSdkCreator(module)?.sdkCreator
+    val sdkCreator = fallback.checkEnvironmentAndPrepareSdkCreator(module)?.getSdkCreator(module)
     if (sdkCreator == null) {
       return false
     }
-    sdkCreator(true).orLogException(logger)
+    sdkCreator.createSdk(needsConfirmation = true).orLogException(logger)
     return true
   }
 

python/python-sdk-configurator/backend/src/impl/ModulesSdkConfigurator.kt

@@ -21,6 +21,8 @@ import com.jetbrains.python.PathShortener
 import com.jetbrains.python.Result
 import com.jetbrains.python.sdk.configuration.CreateSdkInfo
 import com.jetbrains.python.sdk.configuration.PyProjectSdkConfigurationExtension
+import com.jetbrains.python.sdk.configuration.createSdkWithoutConfirmation
+import com.jetbrains.python.sdk.configuration.getSdkCreator
 import com.jetbrains.python.sdk.getOrCreateAdditionalData
 import com.jetbrains.python.sdk.legacy.PythonSdkUtil
 import com.jetbrains.python.sdk.pythonSdk
@@ -141,7 +143,7 @@ internal class ModulesSdkConfigurator private constructor(
           val createInfo = (modules[module.name] ?: error("No create info for module $module, caller broke the contract"))
           when (createInfo) {
             is ModuleCreateInfo.CreateSdkInfoWrapper -> {
-              when (val r = createInfo.createSdkInfo.createSdkWithoutConfirmation()) {
+              when (val r = createInfo.createSdkInfo.createSdkWithoutConfirmation(module)) {
                 is Result.Failure -> { //TODO: Show SDK creation error?
                   logger.warn("Failed to create SDK for ${module.name}: ${r.error}")
                 }
@@ -236,7 +238,7 @@ private suspend fun configureSdkForModuleAutomatically(module: Module, createEnv
 }
 
 private suspend fun CreateSdkInfo.createAndSetToModule(module: Module) {
-  when (val r = sdkCreator(false)) {
+  when (val r = getSdkCreator(module).createSdk(needsConfirmation = false)) {
     is Result.Failure -> {
       logger.trace { "Failed to create sdk for ${module.name} : ${r.error}" }
     }

python/python-sdk/src/com/jetbrains/python/sdk/configuration/CreateSdkInfo.kt

@@ -1,21 +1,38 @@
 package com.jetbrains.python.sdk.configuration
 
 import com.intellij.codeInspection.util.IntentionName
+import com.intellij.openapi.module.Module
 import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.util.NlsSafe
 import com.jetbrains.python.PythonInfo
+import com.jetbrains.python.TraceContext
 import com.jetbrains.python.errorProcessing.PyResult
+import kotlinx.coroutines.withContext
 import org.jetbrains.annotations.ApiStatus
 
-typealias NeedsConfirmation = Boolean
 typealias CheckExistence = Boolean
 typealias CheckToml = Boolean
 typealias EnvExists = Boolean
 
+fun interface SdkCreator {
+  suspend fun createSdk(needsConfirmation: Boolean): PyResult<Sdk?>
+}
+
 @ApiStatus.Internal
-sealed interface CreateSdkInfo : Comparable<CreateSdkInfo> {
+sealed class CreateSdkInfo(private val sdkCreator: SdkCreator) :
+  Comparable<CreateSdkInfo> {
   @get:IntentionName
-  val intentionName: String
-  val sdkCreator: suspend (needsConfirmation: NeedsConfirmation) -> PyResult<Sdk?>
+  abstract val intentionName: String
+
+  /**
+   * Creates SDK for a module named [moduleName]. This does **not** affect the module itself, but just sets user readable title.
+   */
+  fun getSdkCreator(moduleName: @NlsSafe String): SdkCreator = {
+    withContext(TraceContext(moduleName)) {
+      sdkCreator.createSdk(it)
+    }
+  }
+
 
   /**
    * Nullable SDK is only possible when we requested user confirmation but didn't get it. The idea behind this function is to provide
@@ -23,8 +40,11 @@ sealed interface CreateSdkInfo : Comparable<CreateSdkInfo> {
    *
    * It's a temporary solution until we'll be able to remove all custom user dialogs and enable [enableSDKAutoConfigurator] by default.
    * After that lands, we'll get rid of nullable SDK.
+   *
+   * [moduleName] does **not** affect the module itself, but just sets user readable title.
    */
-  suspend fun createSdkWithoutConfirmation(): PyResult<Sdk> = sdkCreator(false).mapSuccess { it!! }
+  suspend fun createSdkWithoutConfirmation(moduleName: @NlsSafe String): PyResult<Sdk> =
+    getSdkCreator(moduleName).createSdk(needsConfirmation = false).mapSuccess { it!! }
 
   /**
    * We want to preserve the initial order, but at the same time existing environment should have a higher priority by default
@@ -35,16 +55,16 @@ sealed interface CreateSdkInfo : Comparable<CreateSdkInfo> {
     return thisExists.compareTo(otherExists)
   }
 
-  data class ExistingEnv(
+  class ExistingEnv internal constructor(
     val pythonInfo: PythonInfo,
     override val intentionName: String,
-    override val sdkCreator: suspend (needsConfirmation: NeedsConfirmation) -> PyResult<Sdk?>,
-  ) : CreateSdkInfo
+    sdkCreator: SdkCreator,
+  ) : CreateSdkInfo(sdkCreator)
 
-  data class WillCreateEnv(
+  class WillCreateEnv internal constructor(
     override val intentionName: String,
-    override val sdkCreator: suspend (needsConfirmation: NeedsConfirmation) -> PyResult<Sdk?>,
-  ) : CreateSdkInfo
+    sdkCreator: SdkCreator,
+  ) : CreateSdkInfo(sdkCreator)
 }
 
 @ApiStatus.Internal
@@ -58,7 +78,7 @@ sealed interface EnvCheckerResult {
 // TODO: Make internal after we drop WSL sdk configurator
 suspend fun prepareSdkCreator(
   envChecker: suspend (CheckExistence) -> EnvCheckerResult,
-  sdkCreator: (EnvExists) -> (suspend (needsConfirmation: NeedsConfirmation) -> PyResult<Sdk?>),
+  sdkCreator: (EnvExists) -> SdkCreator,
 ): CreateSdkInfo? {
   var res = envChecker(true)
   return when (res) {
@@ -78,3 +98,8 @@ suspend fun prepareSdkCreator(
     is EnvCheckerResult.CannotConfigure -> null
   }
 }
+
+fun CreateSdkInfo.getSdkCreator(module: Module): SdkCreator =
+  getSdkCreator(module.name)
+
+suspend fun CreateSdkInfo.createSdkWithoutConfirmation(module: Module): PyResult<Sdk> = createSdkWithoutConfirmation(module.name)

python/python-sdk/tests/com/intellij/python/junit5Tests/unit/ExistingEnvTest.kt

@@ -0,0 +1,30 @@
+package com.intellij.python.junit5Tests.unit
+
+import com.intellij.openapi.util.NlsSafe
+import com.jetbrains.python.PythonInfo
+import com.jetbrains.python.Result
+import com.jetbrains.python.TraceContext
+import com.jetbrains.python.isFailure
+import com.jetbrains.python.psi.LanguageLevel
+import com.jetbrains.python.sdk.configuration.CreateSdkInfo
+import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.runBlocking
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.Test
+
+internal class ExistingEnvTest {
+  private companion object {
+    fun ExistingEnv(expectedTitle: @NlsSafe String) =
+      CreateSdkInfo.ExistingEnv(PythonInfo(LanguageLevel.PYTHON31), "...") {
+        Assertions.assertEquals(expectedTitle, currentCoroutineContext()[TraceContext]?.title)
+        Result.localizedError("...")
+      }
+  }
+
+  @Test
+  fun testContext(): Unit = runBlocking {
+    val title = "some project"
+    val result = ExistingEnv(expectedTitle = title).getSdkCreator(title).createSdk(needsConfirmation = false)
+    Assertions.assertTrue(result.isFailure)
+  }
+}

python/src/com/jetbrains/python/inspections/interpreter/PyInterpreterInspection.kt

@@ -38,6 +38,7 @@ import com.jetbrains.python.sdk.conda.PyCondaSdkCustomizer
 import com.jetbrains.python.sdk.configuration.CreateSdkInfo
 import com.jetbrains.python.sdk.configuration.CreateSdkInfoWithTool
 import com.jetbrains.python.sdk.configuration.PyProjectSdkConfiguration
+import com.jetbrains.python.sdk.configuration.createSdkWithoutConfirmation
 import com.jetbrains.python.sdk.legacy.PythonSdkUtil
 import com.jetbrains.python.sdk.service.PySdkService.Companion.pySdkService
 import com.jetbrains.python.ui.PyUiUtil
@@ -271,7 +272,7 @@ private suspend fun Module.getQuickFixBySdkSuggestion(i: ModuleCreateInfo?): Fin
     when (val createSdkInfo = i.createSdkInfo) {
       is CreateSdkInfo.ExistingEnv -> {
         logger.trace { "$this: Files already exist, just create sn SDK" }
-        when (val creationResult = createSdkInfo.createSdkWithoutConfirmation()) {
+        when (val creationResult = createSdkInfo.createSdkWithoutConfirmation(module = this)) {
           is Result.Failure -> {
             logger.warn("Can't create SDK for $this : ${creationResult.error}")
             FindQuickFixResult.NoSuggestion

python/src/com/jetbrains/python/sdk/configuration/PyProjectSdkConfiguration.kt

@@ -43,11 +43,11 @@ object PyProjectSdkConfiguration {
   }
 
   suspend fun setSdkUsingCreateSdkInfo(
-    module: Module, createSdkInfoWithTool: CreateSdkInfoWithTool, needsConfirmation: NeedsConfirmation,
+    module: Module, createSdkInfoWithTool: CreateSdkInfoWithTool, needsConfirmation: Boolean,
   ): Boolean = withContext(Dispatchers.Default) {
     thisLogger().debug("Configuring sdk using ${createSdkInfoWithTool.toolId}")
 
-    val sdk = createSdkInfoWithTool.createSdkInfo.sdkCreator(needsConfirmation).getOr {
+    val sdk = createSdkInfoWithTool.createSdkInfo.getSdkCreator(module).createSdk(needsConfirmation).getOr {
       ShowingMessageErrorSync.emit(it.error, module.project)
       return@withContext true
     } ?: return@withContext false