Python: move classes to access pythons created by setup-env script into right module.

We now store classes to access pythons installed by tests setup env right next to the script

GitOrigin-RevId: 922e5cfbb89e1e550679796e02deb24553ad6852

python/setup-test-environment/README.txt

@@ -0,0 +1,3 @@
+Python tests require pythons with certain modules.
+
+Run `build.gradle.kts` to install them. Then, use classes from this module to access them.

python/setup-test-environment/intellij.python.community.testFramework.testEnv.iml

@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="JAVA_MODULE" version="4">
+  <component name="NewModuleRootManager" inherit-compiler-output="true">
+    <exclude-output />
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+    <orderEntry type="library" name="kotlin-stdlib" level="project" />
+    <orderEntry type="module" module-name="intellij.platform.util" />
+    <orderEntry type="module" module-name="intellij.python.community" />
+    <orderEntry type="module" module-name="intellij.platform.testFramework" />
+    <orderEntry type="module" module-name="intellij.python.sdk" />
+    <orderEntry type="module" module-name="intellij.python.community.impl" />
+    <orderEntry type="module" module-name="intellij.platform.execution" />
+    <orderEntry type="library" name="kotlin-reflect" level="project" />
+  </component>
+</module>

python/setup-test-environment/src/com/intellij/python/community/testFramework/testEnv/PyEnvTestSettings.kt

@@ -0,0 +1,130 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.python.community.testFramework.testEnv
+
+import com.intellij.python.community.testFramework.testEnv.PyEnvTestSettings.Companion.PATH_TO_TEST_ENV_PYTHON_INTERPRETERS
+import com.intellij.testFramework.UsefulTestCase.IS_UNDER_TEAMCITY
+import com.intellij.util.SystemProperties
+import org.jetbrains.annotations.ApiStatus
+import org.jetbrains.annotations.TestOnly
+import java.io.File
+
+
+/**
+ * Configures env test environment using env vars and properties.
+ * Environment variables are also used in gradle script (setup-test-environment)
+ */
+@TestOnly
+@ApiStatus.Internal
+data class PyEnvTestSettings(
+  private val folderWithCPythons: String?,
+  private val folderWithCondas: String?,
+  private val additionalInterpreters: List<File>,
+  @get:JvmName("useRemoteSdk")
+  val useRemoteSdk: Boolean,
+  val isEnvConfiguration: Boolean,
+  val isUnderTeamCity: Boolean,
+) {
+  private val foldersWithPythons: List<File> = listOfNotNull(folderWithCPythons, folderWithCondas).map { File(it) }
+
+  /**
+   * Paths to all existing python SDKs (as directories with python installation)
+   */
+  @TestOnly
+  @get:JvmName("getPythons")
+  internal val pythons: List<File> = foldersWithPythons
+    .filter(File::exists)
+    .flatMap { it.listFiles()?.toList() ?: emptyList() }
+    .plus(additionalInterpreters)
+    .filter { it.isDirectory && !it.name.startsWith('.') }
+
+  /**
+   * Configuration in readable format
+   */
+  fun reportConfiguration() = (PyTestEnvVars.getEnvValues() + listOf(toString())).joinToString("\n")
+
+  companion object {
+    // If you decided to change the path, make sure to update it in community/python/setup-test-environment/build.gradle.kts
+    private const val PATH_TO_TEST_ENV_PYTHON_INTERPRETERS = "community/python/setup-test-environment/build/pythons"
+
+    /**
+     * Tries to resolve [PATH_TO_TEST_ENV_PYTHON_INTERPRETERS] folder from current working dir.
+     * It allows automatically detecting the local folder with interpreters built by [community/python/setup-test-environment/build.gradle.kts]
+     */
+    private fun detectDefaultPyInterpretersFolders(): List<File> {
+      var currentFile: File? = File(System.getProperty("user.dir"))
+      while (currentFile != null) {
+        val pythonInterpretersFolder = currentFile.resolve(PATH_TO_TEST_ENV_PYTHON_INTERPRETERS)
+        if (pythonInterpretersFolder.exists()) {
+          return pythonInterpretersFolder.listFiles { it.isDirectory && !it.name.startsWith('.') }.toList()
+        }
+        currentFile = currentFile.parentFile
+
+      }
+      return emptyList()
+    }
+
+    fun fromEnvVariables(): PyEnvTestSettings {
+      val isUnderTeamCity = IS_UNDER_TEAMCITY
+      return PyEnvTestSettings(
+        folderWithCPythons = PyTestEnvVars.PYCHARM_PYTHONS.getValue(),
+        folderWithCondas = PyTestEnvVars.PYCHARM_PYTHON_VIRTUAL_ENVS.getValue(),
+        useRemoteSdk = SystemProperties.getBooleanProperty("pycharm.run_remote", false) || PyTestEnvVars.PYCHARM_RUN_REMOTE.isSet(),
+        isEnvConfiguration = System.getProperty("pycharm.env") != null
+                             || PyTestEnvVars.PYCHARM_ENV.toString() in System.getenv(),
+        isUnderTeamCity = isUnderTeamCity,
+        additionalInterpreters = PyTestEnvVars.PYCHARM_PYTHON_ENVS.getValue()?.split(File.pathSeparator)
+                                   ?.map { File(it) }
+                                   ?.toList()
+                                 ?: if (isUnderTeamCity) {
+                                   emptyList()
+                                 }
+                                 else {
+                                   detectDefaultPyInterpretersFolders()
+                                 },
+      )
+    }
+  }
+}
+
+/**
+ * Env variables used to configure tests
+ */
+@ApiStatus.Internal
+enum class PyTestEnvVars(private val getVarName: (PyTestEnvVars) -> String = { it.name }) {
+  /**
+   * Path to folder with CPythons
+   */
+  PYCHARM_PYTHONS,
+
+  /**
+   * Path to folder with condas
+   */
+  PYCHARM_PYTHON_VIRTUAL_ENVS,
+
+  /**
+   * [File.separator] separated list of full paths to pythons (including binary) to add to folders found in [PYCHARM_PYTHON_VIRTUAL_ENVS]
+   * and [PYCHARM_PYTHONS]
+   */
+  PYCHARM_PYTHON_ENVS,
+
+  /**
+   * Set if launched using "PyEnvTests"
+   */
+  PYCHARM_ENV,
+
+  /**
+   * Only run remote-based tests
+   */
+  PYCHARM_RUN_REMOTE;
+
+
+  companion object {
+    fun getEnvValues() = entries.map { "$it : ${it.getValue()}" }
+  }
+
+  override fun toString() = getVarName(this)
+
+  fun getValue(): String? = System.getenv(toString())
+
+  fun isSet() = getValue() != null
+}

python/setup-test-environment/src/com/intellij/python/community/testFramework/testEnv/PythonType.kt

@@ -0,0 +1,136 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.python.community.testFramework.testEnv
+
+import com.intellij.execution.processTools.getResultStdout
+import com.intellij.execution.target.local.LocalTargetEnvironmentRequest
+import com.intellij.openapi.diagnostic.logger
+import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.vfs.newvfs.impl.VfsRootAccess
+import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
+import com.jetbrains.python.PythonBinary
+import com.jetbrains.python.packaging.findCondaExecutableRelativeToEnv
+import com.jetbrains.python.sdk.conda.TargetEnvironmentRequestCommandExecutor
+import com.jetbrains.python.sdk.flavors.PythonSdkFlavor
+import com.jetbrains.python.sdk.flavors.conda.PyCondaEnv
+import com.jetbrains.python.sdk.flavors.conda.PyCondaEnvIdentity
+import com.jetbrains.python.venvReader.VirtualEnvReader
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.runBlocking
+import org.jetbrains.annotations.NonNls
+import java.nio.file.Path
+import java.util.concurrent.ConcurrentHashMap
+import kotlin.Result.Companion.failure
+
+/**
+ * Gradle script installs two types of python: conda and vanilla. Env could be obtained by [getTestEnvironment] which also provides closable
+ * to clean up after the usage
+ */
+sealed class PythonType<T : Any>(private val tag: @NonNls String? = null) {
+  private companion object {
+    val cache: MutableMap<Set<String>, List<Path>> = ConcurrentHashMap()
+  }
+
+  /**
+   * Returns all test environments ordered from newest (highest) to oldest: each must be closed after the test.
+   * If in doubt, take first
+   */
+  suspend fun getTestEnvironments(vararg additionalTags: @NonNls String): Flow<Pair<T, AutoCloseable>> {
+    val key = setOf(*additionalTags)
+    val pythons = cache.getOrPut(key) {
+      PyEnvTestSettings
+        .fromEnvVariables()
+        .pythons
+        .map { it.toPath() }
+        .mapNotNull { path ->
+          val binary = VirtualEnvReader.Instance.findPythonInPythonRoot(path)
+                       ?: error("No python in $path")
+          val flavor = PythonSdkFlavor.tryDetectFlavorByLocalPath(binary.toString())
+                       ?: error("Unknown flavor: $binary")
+          flavor.getVersionString(binary.toString())?.let { path to PythonSdkFlavor.getLanguageLevelFromVersionStringStatic(it) }
+          ?: error("Can't get language level for $flavor , $binary")
+        }
+        .sortedByDescending { (_, languageLevel) -> languageLevel }
+        .filter { (_, languageLevel) -> languageLevel.isPy3K }
+        .map { (path, _) -> path }
+    }
+    return pythons.asFlow()
+      .filter { typeMatchesEnv(it, *additionalTags) }
+      .map { envDir ->
+        pythonPathToEnvironment(
+          VirtualEnvReader.Instance.findPythonInPythonRoot(envDir)
+          ?: error("Can't find python binary in $envDir"), envDir) // This is a misconfiguration, hence an error
+      }
+  }
+
+
+  /**
+   * Returns first (whatever it means) test environment and closable that must be closed after the test
+   */
+  suspend fun getTestEnvironment(vararg additionalTags: @NonNls String): Result<Pair<T, AutoCloseable>> =
+    getTestEnvironments(*additionalTags).firstOrNull()?.let { Result.success(it) }
+    ?: failure(AssertionError("No python found. See ${PyEnvTestSettings::class} class for more info"))
+
+
+  protected abstract suspend fun pythonPathToEnvironment(pythonBinary: PythonBinary, envDir: Path): Pair<T, AutoCloseable>
+
+
+  data object VanillaPython3 : PythonType<PythonBinary>() {
+    // Python is directly executable
+    override suspend fun pythonPathToEnvironment(pythonBinary: PythonBinary, envDir: Path): Pair<PythonBinary, AutoCloseable> {
+      val disposable = Disposer.newDisposable("Python tests disposable for VfsRootAccess")
+      // We might have python installation outside the project root, but we still need to have access to it.
+      VfsRootAccess.allowRootAccess(disposable, pythonBinary.parent.toString())
+      return Pair(pythonBinary, AutoCloseable {
+        Disposer.dispose(disposable)
+      })
+    }
+  }
+
+
+  data object Conda : PythonType<PyCondaEnv>("conda") {
+
+    override suspend fun pythonPathToEnvironment(pythonBinary: PythonBinary, envDir: Path): Pair<PyCondaEnv, AutoCloseable> {
+      // First, find python binary, then calculate conda from it as env stores "conda" as a regular env
+      val condaPath = findCondaExecutableRelativeToEnv(pythonBinary) ?: error("Conda root $pythonBinary doesn't have conda binary")
+
+      // We'll remove then on close
+      val condaEnvsBeforeTest = getCondaNames(condaPath)
+
+      val cleanupCondas = AutoCloseable {
+        runBlocking(Dispatchers.IO) {
+          val condasToRemove = getCondaNames(condaPath)
+          condasToRemove.removeAll(condaEnvsBeforeTest)
+          for (envName in condasToRemove) {
+            println("Removing $envName")
+
+            for (arg in arrayOf("--name", "-p")) {
+              val args = arrayOf(condaPath.toString(), "remove", arg, envName, "--all", "-y")
+              Runtime.getRuntime().exec(args).getResultStdout().getOrElse {
+                logger<Conda>().warn(it)
+              }
+            }
+          }
+        }
+      }
+      return Pair(PyCondaEnv(PyCondaEnvIdentity.UnnamedEnv(envDir.toString(), isBase = true), condaPath.toString()), cleanupCondas)
+    }
+
+    private suspend fun getCondaNames(condaPath: Path) =
+      PyCondaEnv.getEnvs(TargetEnvironmentRequestCommandExecutor(LocalTargetEnvironmentRequest()),
+                         condaPath.toString()).getOrThrow()
+        .map { it.envIdentity.userReadableName }
+        .toMutableSet()
+  }
+
+
+  @RequiresBackgroundThread
+  private fun typeMatchesEnv(env: Path, vararg additionalTags: @NonNls String): Boolean {
+    val envTags = loadEnvTags(env)
+
+    for (badTag in PythonType::class.sealedSubclasses.filterNot { it.isInstance(this) }.map { it.objectInstance!!.tag }) {
+      if (badTag != null && badTag in envTags) return false
+    }
+    return (tag == null || tag in envTags) && additionalTags.all { it in envTags }
+  }
+}

python/setup-test-environment/src/com/intellij/python/community/testFramework/testEnv/envTags.kt

@@ -0,0 +1,29 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.python.community.testFramework.testEnv
+
+import com.intellij.openapi.diagnostic.fileLogger
+import com.jetbrains.python.venvReader.Directory
+import java.io.IOException
+import java.nio.file.Path
+import kotlin.io.path.isDirectory
+import kotlin.io.path.readLines
+
+private const val TAGS_FILE = "tags.txt"
+
+/**
+ * Loads [TAGS_FILE] from [pythonEnv] which should be a path to a python installation (either binary or directory) created by
+ * `community/python/setup-test-environment`.
+ *
+ * Returns set of tags i.e. `django` if interpreter has django.
+ */
+fun loadEnvTags(pythonEnv: Path): Set<String> {
+  val envDir: Directory = if (pythonEnv.isDirectory()) pythonEnv else pythonEnv.parent
+  val tagsFile = envDir.resolve(TAGS_FILE)
+  try {
+    return tagsFile.readLines().toSet()
+  }
+  catch (_: IOException) {
+    fileLogger().warn("Can't read $tagsFile")
+    return emptySet()
+  }
+}

python/setup-test-environment/src/com/intellij/python/community/testFramework/testEnv/package-info.java

@@ -0,0 +1,8 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+/**
+ * Package to access pythons created by <code>community/python/setup-test-environment/build.gradle.kts</code>
+ */
+@ApiStatus.Internal
+package com.intellij.python.community.testFramework.testEnv;
+
+import org.jetbrains.annotations.ApiStatus;