[python][hatch] support matrix environments (PY-60410)

+ add HatchServiceTest

GitOrigin-RevId: a19006a801ca5848ec3bc5a0398c17e4516c10bb

python/python-hatch/src/com/intellij/python/hatch/cli/AsciiTable.kt

@@ -0,0 +1,47 @@
+package com.intellij.python.hatch.cli
+
+/**
+ * Simple table without joined cells.
+ * Each row has the same number of cells as the header.
+ */
+internal data class AsciiTable(val headers: List<String>, val rows: List<List<String>>) {
+  fun findColumnIdx(caption: String): Int? = headers.indexOf(caption).takeIf { it >= 0 }
+  fun List<String>.cell(columnIdx: Int?): String? = columnIdx?.let { this[it] }
+}
+
+
+/**
+ * Parses Tables like this:
+ *
+ * +------+------+
+ * |  H1  |  H2  |
+ * +======+======+
+ * | R1_1 | R1_2 |
+ * +------+------+
+ * | R2_1 | R2_2 |
+ * +------+------+
+ * ...
+ * +------+------+
+ */
+internal fun String.parseAsciiTable(): AsciiTable? {
+  val lines = this.trim().lines()
+  val columns = lines.first().count { it == '+' } - 1
+  if (columns <= 0) return null
+
+  val data = buildList {
+    val row = Array(columns) { "" }
+    for (line in lines.drop(1)) {
+      if (line.startsWith('+')) {
+        add(row.map { it.trim() })
+        row.fill("")
+        continue
+      }
+
+      val cells = line.splitToSequence('|').map { it.trim() }.toList()
+      for (col in 0..<columns) {
+        row[col] += "\n${cells[col + 1]}"
+      }
+    }
+  }
+  return AsciiTable(headers = data.first(), rows = data.drop(1))
+}

python/python-hatch/src/com/intellij/python/hatch/cli/HatchEnv.kt

@@ -231,27 +231,24 @@ class HatchEnv(runtime: HatchRuntime) : HatchCommand("env", runtime) {
     val expectedOutput = """^\s+Standalone\s*\n((?:[+|].*[+|]\n)+)(?:\s+Matrices\s*\n((?:[+|].*[+|]\n)+))?$""".toRegex()
 
     return executeAndMatch("show", "--ascii", *options, *envs, expectedOutput = expectedOutput) { matchResult ->
-      val tables = buildMap {
-        val (standaloneTable, matricesTable) = matchResult.destructured
-        put(HatchEnvironmentType.STANDALONE, standaloneTable.parseHatchEnvironments())
-        put(HatchEnvironmentType.MATRICES, matricesTable.parseHatchEnvironments())
-      }
-      Result.success(tables)
+      val (standaloneTable, matricesTable) = matchResult.destructured
+      val standalone = standaloneTable.parseAsciiTable()?.parseHatchEnvironments()?.map { it.first } ?: emptyList()
+      val matrices = matricesTable.parseAsciiTable()?.parseHatchEnvironments()?.mapNotNull {
+        it.second?.let { envs -> HatchMatrixEnvironment(it.first, envs) }
+      } ?: emptyList()
+      Result.success(HatchEnvironments(standalone, matrices))
     }
   }
 }
 
-typealias HatchEnvironments = Map<HatchEnvironmentType, List<HatchEnvironment>>
-
-enum class HatchEnvironmentType {
-  STANDALONE,
-  MATRICES,
-}
+data class HatchEnvironments(
+  val standalone: List<HatchEnvironment>,
+  val matrices: List<HatchMatrixEnvironment>,
+)
 
 data class HatchEnvironment(
   val name: @NlsSafe String,
   val type: @NlsSafe String,
-  val envs: String? = null,
   val features: String? = null,
   val dependencies: String? = null,
   val environmentVariables: String? = null,
@@ -263,56 +260,32 @@ data class HatchEnvironment(
   }
 }
 
-private fun String.parseHatchEnvironments(): List<HatchEnvironment> {
-  val table = parseTable() ?: return emptyList()
-  val nameIdx = table.findColumnIdx("Name") ?: error("Name column not found")
-  val typeIdx = table.findColumnIdx("Type") ?: error("Type column not found")
-  val envIdx = table.findColumnIdx("Env")
-  val featuresIdx = table.findColumnIdx("Features")
-  val dependenciesIdx = table.findColumnIdx("Features")
-  val environmentVariablesIdx = table.findColumnIdx("EnvironmentVariables")
-  val scriptsIdx = table.findColumnIdx("Scripts")
-  val descriptionIdx = table.findColumnIdx("Description")
+data class HatchMatrixEnvironment(
+  val hatchEnvironment: HatchEnvironment,
+  val envs: List<String>,
+)
 
-  return table.rows.map { row ->
+
+private fun AsciiTable.parseHatchEnvironments(): List<Pair<HatchEnvironment, List<String>?>> {
+  val nameIdx = findColumnIdx("Name") ?: error("Name column not found")
+  val typeIdx = findColumnIdx("Type") ?: error("Type column not found")
+  val featuresIdx = findColumnIdx("Features")
+  val dependenciesIdx = findColumnIdx("Dependencies")
+  val environmentVariablesIdx = findColumnIdx("Environment variables")
+  val scriptsIdx = findColumnIdx("Scripts")
+  val descriptionIdx = findColumnIdx("Description")
+  val envsIdx = findColumnIdx("Envs")
+
+  return rows.map { row ->
+    val matrixEnvironments = envsIdx?.let { idx -> row[idx].lines().map { it.trim() } }
     HatchEnvironment(
       name = row[nameIdx],
       type = row[typeIdx],
-      envs = envIdx?.let { row[it] },
-      features = featuresIdx?.let { row[it] },
-      dependencies = dependenciesIdx?.let { row[it] },
-      environmentVariables = environmentVariablesIdx?.let { row[it] },
-      scripts = scriptsIdx?.let { row[it] },
-      description = descriptionIdx?.let { row[it] },
-    )
+      features = row.cell(featuresIdx),
+      dependencies = row.cell(dependenciesIdx),
+      environmentVariables = row.cell(environmentVariablesIdx),
+      scripts = row.cell(scriptsIdx),
+      description = row.cell(descriptionIdx),
+    ) to matrixEnvironments
   }
 }
-
-
-private data class Table(val headers: List<String>, val rows: List<List<String>>)
-
-private fun Table.findColumnIdx(name: String): Int? = headers.indexOf(name).takeIf { it >= 0 }
-
-
-private fun String.parseTable(): Table? {
-  val lines = this.trim().lines()
-  val columns = lines.first().count { it == '+' } - 1
-  if (columns <= 0) return null
-
-  val data = buildList {
-    val currentRow = Array(columns) { "" }
-    for (line in lines.drop(1)) {
-      if (line.startsWith('+')) {
-        add(currentRow.map { it.trim() })
-        currentRow.fill("")
-        continue
-      }
-
-      val cells = line.splitToSequence('|').map { it.trim() }.toList()
-      for (col in 0..<columns) {
-        currentRow[col] += "\n${cells[col + 1]}"
-      }
-    }
-  }
-  return Table(headers = data.first(), rows = data.drop(1))
-}

python/python-hatch/src/com/intellij/python/hatch/hatch.kt

@@ -41,13 +41,13 @@ class FileSystemOperationHatchError(eelFsError: EelFsError) : HatchError(
 )
 
 
-data class HatchStandaloneEnvironment(
+data class HatchVirtualEnvironment(
   val hatchEnvironment: HatchEnvironment,
   val pythonVirtualEnvironment: PythonVirtualEnvironment,
 ) {
   companion object {
-    val AVAILABLE_ENVIRONMENTS_FOR_NEW_PROJECT: List<HatchStandaloneEnvironment> = listOf(
-      HatchStandaloneEnvironment(HatchEnvironment.DEFAULT, PythonVirtualEnvironment.NotExisting())
+    val AVAILABLE_ENVIRONMENTS_FOR_NEW_PROJECT: List<HatchVirtualEnvironment> = listOf(
+      HatchVirtualEnvironment(HatchEnvironment.DEFAULT, PythonVirtualEnvironment.NotExisting())
     )
   }
 }
@@ -59,6 +59,11 @@ sealed interface PythonVirtualEnvironment {
   data class NotExisting(override val pythonHomePath: PythonHomePath? = null) : PythonVirtualEnvironment
 }
 
+data class ProjectStructure(
+  val sourceRoot: Path?,
+  val testRoot: Path?,
+)
+
 interface HatchService {
   fun getWorkingDirectoryPath(): Path
 
@@ -66,7 +71,7 @@ interface HatchService {
   suspend fun isHatchManagedProject(): Result<Boolean, PyError>
 
   @RequiresBackgroundThread
-  suspend fun createNewProject(projectName: String): Result<Unit, PyError>
+  suspend fun createNewProject(projectName: String): Result<ProjectStructure, PyError>
 
   /**
    * param[basePythonBinaryPath] base python for environment, the one on the PATH should be used if null.
@@ -76,7 +81,7 @@ interface HatchService {
   suspend fun createVirtualEnvironment(basePythonBinaryPath: PythonBinary? = null, envName: String? = null): Result<PythonVirtualEnvironment.Existing, PyError>
 
   @RequiresBackgroundThread
-  suspend fun findStandaloneEnvironments(): Result<List<HatchStandaloneEnvironment>, PyError>
+  suspend fun findVirtualEnvironments(): Result<List<HatchVirtualEnvironment>, PyError>
 }
 
 /**

python/python-hatch/src/com/intellij/python/hatch/service/CliBasedHatchService.kt

@@ -1,6 +1,8 @@
 package com.intellij.python.hatch.service
 
 import com.intellij.platform.eel.fs.EelFileSystemApi
+import com.intellij.platform.eel.fs.EelFileSystemApi.ReplaceExistingDuringMove.DO_NOT_REPLACE_DIRECTORIES
+import com.intellij.platform.eel.fs.move
 import com.intellij.platform.eel.getOr
 import com.intellij.platform.eel.provider.asEelPath
 import com.intellij.platform.eel.provider.asNioPath
@@ -8,7 +10,9 @@ import com.intellij.platform.eel.provider.getEelDescriptor
 import com.intellij.python.community.execService.ExecService
 import com.intellij.python.community.execService.WhatToExec.Binary
 import com.intellij.python.hatch.*
-import com.intellij.python.hatch.cli.HatchEnvironmentType
+import com.intellij.python.hatch.cli.ENV_TYPE_VIRTUAL
+import com.intellij.python.hatch.cli.HatchEnvironment
+import com.intellij.python.hatch.cli.HatchEnvironments
 import com.intellij.python.hatch.runtime.HatchRuntime
 import com.intellij.python.hatch.runtime.createHatchRuntime
 import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
@@ -17,6 +21,9 @@ import com.jetbrains.python.PythonHomePath
 import com.jetbrains.python.Result
 import com.jetbrains.python.errorProcessing.PyError
 import com.jetbrains.python.resolvePythonBinary
+import kotlinx.coroutines.async
+import kotlinx.coroutines.awaitAll
+import kotlinx.coroutines.coroutineScope
 import java.nio.file.Path
 import kotlin.io.path.exists
 import kotlin.io.path.isDirectory
@@ -30,8 +37,8 @@ internal class CliBasedHatchService private constructor(
   companion object {
     suspend operator fun invoke(workingDirectoryPath: Path, hatchExecutablePath: Path?): Result<CliBasedHatchService, PyError> {
       val hatchRuntime = createHatchRuntime(
-        hatchExecutablePath=hatchExecutablePath,
-        workingDirectoryPath=workingDirectoryPath,
+        hatchExecutablePath = hatchExecutablePath,
+        workingDirectoryPath = workingDirectoryPath,
       ).getOr { return it }
       return Result.success(CliBasedHatchService(workingDirectoryPath, hatchRuntime))
     }
@@ -52,42 +59,46 @@ internal class CliBasedHatchService private constructor(
     return Result.success(isHatchManaged)
   }
 
-  @RequiresBackgroundThread
-  override suspend fun findStandaloneEnvironments(): Result<List<HatchStandaloneEnvironment>, PyError> {
-    val hatchEnv = hatchRuntime.hatchCli().env()
-    val environments = hatchEnv.show().getOr { return it }
-    val standaloneEnvironments = environments.getOrDefault(HatchEnvironmentType.STANDALONE, emptyList())
+  suspend fun <A, B> Iterable<A>.parallelMap(f: suspend (A) -> B): List<B> = coroutineScope {
+    map { async { f(it) } }.awaitAll()
+  }
 
-    val available = standaloneEnvironments.mapNotNull { env ->
-      val pythonHomePath = hatchEnv.find(env.name).getOr { return@mapNotNull null } ?: return@mapNotNull null
-      val pythonVirtualEnvironment = pythonHomePath.toPythonVirtualEnvironment().getOr { return@mapNotNull null }
-      HatchStandaloneEnvironment(
+  @RequiresBackgroundThread
+  override suspend fun findVirtualEnvironments(): Result<List<HatchVirtualEnvironment>, PyError> {
+    val hatchEnv = hatchRuntime.hatchCli().env()
+    val environments: HatchEnvironments = hatchEnv.show().getOr { return it }
+    val virtualEnvironments = environments.getAvailableVirtualHatchEnvironments()
+
+    val available = virtualEnvironments.parallelMap { env ->
+      val pythonHomePath = hatchEnv.find(env.name).getOr { return@parallelMap null } ?: return@parallelMap null
+      val pythonVirtualEnvironment = pythonHomePath.toPythonVirtualEnvironment().getOr { return@parallelMap null }
+      HatchVirtualEnvironment(
         hatchEnvironment = env,
         pythonVirtualEnvironment = pythonVirtualEnvironment
       )
-    }
+    }.filterNotNull()
 
     return Result.success(available)
   }
 
+
   @RequiresBackgroundThread
-  override suspend fun createNewProject(projectName: String): Result<Unit, PyError> {
+  override suspend fun createNewProject(projectName: String): Result<ProjectStructure, PyError> {
     val eelApi = workingDirectoryPath.getEelDescriptor().upgrade()
     val tempDir = eelApi.fs.createTemporaryDirectory(EelFileSystemApi.CreateTemporaryEntryOptions.Builder().build()).getOr { failure ->
       return Result.failure(FileSystemOperationHatchError(failure.error))
     }
 
     hatchRuntime.hatchCli().new(projectName, tempDir.asNioPath()).getOr { return it }
-    eelApi.fs.move(
-      tempDir,
-      workingDirectoryPath.asEelPath(),
-      EelFileSystemApi.ReplaceExistingDuringMove.DO_NOT_REPLACE_DIRECTORIES,
-      true
-    ).getOr { failure ->
+    val target = workingDirectoryPath.asEelPath()
+    eelApi.fs.move(tempDir, target).replaceExisting(DO_NOT_REPLACE_DIRECTORIES).eelIt().getOr { failure ->
       return Result.failure(FileSystemOperationHatchError(failure.error))
     }
 
-    return Result.success(Unit)
+    return Result.success(ProjectStructure(
+      sourceRoot = target.asNioPath().resolve("src").takeIf { it.isDirectory() },
+      testRoot = target.asNioPath().resolve("tests").takeIf { it.isDirectory() },
+    ))
   }
 
   @RequiresBackgroundThread
@@ -120,4 +131,22 @@ internal suspend fun PythonHomePath.toPythonVirtualEnvironment(): Result<PythonV
     else -> PythonVirtualEnvironment.Existing(this, pythonVersion)
   }
   return Result.success(pythonVirtualEnvironment)
+}
+
+private fun HatchEnvironments.getAvailableVirtualHatchEnvironments(): List<HatchEnvironment> {
+  val matricesFlatted = matrices.flatMap { matrixEnvironment ->
+    matrixEnvironment.envs.map { envName ->
+      with(matrixEnvironment.hatchEnvironment) {
+        HatchEnvironment(
+          name = envName,
+          type = type,
+          dependencies = dependencies,
+          environmentVariables = environmentVariables,
+          scripts = scripts,
+          description = description,
+        )
+      }
+    }
+  }
+  return (standalone + matricesFlatted).filter { it.type == ENV_TYPE_VIRTUAL }
 }