feat(pypi): attempt more python3 candidates (#1608)

2 files changed, 66 insertions, 12 deletions

diff --git a/lua/mason-core/installer/managers/pypi.lua b/lua/mason-core/installer/managers/pypi.lua
index 7c0e5bb9..bda8925e 100644
--- a/lua/mason-core/installer/managers/pypi.lua
+++ b/lua/mason-core/installer/managers/pypi.lua
@@ -6,18 +6,71 @@ local installer = require "mason-core.installer"
 local log = require "mason-core.log"
 local path = require "mason-core.path"
 local platform = require "mason-core.platform"
+local semver = require "mason-core.semver"
+local spawn = require "mason-core.spawn"
 
 local M = {}
 
 local VENV_DIR = "venv"
 
+local is_executable = _.compose(_.equals(1), vim.fn.executable)
+
+---@async
+---@param candidates string[]
+local function resolve_python3(candidates)
+    a.scheduler()
+    local available_candidates = _.filter(is_executable, candidates)
+    for __, candidate in ipairs(available_candidates) do
+        ---@type string
+        local version_output = spawn[candidate]({ "--version" }):map(_.prop "stdout"):get_or_else ""
+        local ok, version = pcall(semver.new, version_output:match "Python (3%.%d+.%d+)")
+        if ok then
+            return { executable = candidate, version = version }
+        end
+    end
+    return nil
+end
+
+---@param min_version? Semver
+local function get_versioned_candidates(min_version)
+    return _.filter_map(function(pair)
+        local version, executable = unpack(pair)
+        if not min_version or version > min_version then
+            return Optional.of(executable)
+        else
+            return Optional.empty()
+        end
+    end, {
+        { semver.new "3.12.0", "python3.12" },
+        { semver.new "3.11.0", "python3.11" },
+        { semver.new "3.10.0", "python3.10" },
+        { semver.new "3.9.0", "python3.9" },
+        { semver.new "3.8.0", "python3.8" },
+        { semver.new "3.7.0", "python3.7" },
+        { semver.new "3.6.0", "python3.6" },
+    })
+end
+
 ---@async
----@param py_executables string[]
-local function create_venv(py_executables)
+local function create_venv()
+    local stock_candidates = platform.is.win and { "python", "python3" } or { "python3", "python" }
+    local stock_target = resolve_python3(stock_candidates)
+    local _ = stock_target and log.fmt_debug("Resolved stock python3 installation version %s", stock_target.version)
+    local versioned_candidates = get_versioned_candidates(stock_target and stock_target.version)
+    log.debug("Resolving versioned python3 candidates", versioned_candidates)
+    local target = resolve_python3(versioned_candidates) or stock_target
     local ctx = installer.context()
-    return Optional.of_nilable(_.find_first(function(executable)
-        return ctx.spawn[executable]({ "-m", "venv", VENV_DIR }):is_success()
-    end, py_executables)):ok_or "Failed to create python3 virtual environment."
+    if not target then
+        ctx.stdio_sink.stderr(
+            ("Unable to find python3 installation. Tried the following candidates: %s.\n"):format(
+                _.join(", ", _.concat(stock_candidates, versioned_candidates))
+            )
+        )
+        return Result.failure "Failed to find python3 installation."
+    end
+    log.fmt_debug("Found python3 installation version=%s, executable=%s", target.version, target.executable)
+    ctx.stdio_sink.stdout "Creating virtual environment…\n"
+    return ctx.spawn[target.executable] { "-m", "venv", VENV_DIR }
 end
 
 ---@param ctx InstallContext
@@ -70,15 +123,9 @@ function M.init(opts)
         log.fmt_debug("pypi: init", opts)
         local ctx = installer.context()
 
-        a.scheduler()
-
-        local executables = platform.is.win and { "python", "python3" } or { "python3", "python" }
-
         -- pip3 will hardcode the full path to venv executables, so we need to promote cwd to make sure pip uses the final destination path.
         ctx:promote_cwd()
-
-        ctx.stdio_sink.stdout "Creating virtual environment…\n"
-        try(create_venv(executables))
+        try(create_venv())
 
         if opts.upgrade_pip then
             ctx.stdio_sink.stdout "Upgrading pip inside the virtual environment…\n"
diff --git a/tests/mason-core/installer/managers/pypi_spec.lua b/tests/mason-core/installer/managers/pypi_spec.lua
index 353606aa..4eff4092 100644
--- a/tests/mason-core/installer/managers/pypi_spec.lua
+++ b/tests/mason-core/installer/managers/pypi_spec.lua
@@ -1,7 +1,9 @@
+local Result = require "mason-core.result"
 local installer = require "mason-core.installer"
 local match = require "luassert.match"
 local path = require "mason-core.path"
 local pypi = require "mason-core.installer.managers.pypi"
+local spawn = require "mason-core.spawn"
 local spy = require "luassert.spy"
 local stub = require "luassert.stub"
 
@@ -16,6 +18,11 @@ local function venv_py(ctx)
 end
 
 describe("pypi manager", function()
+    before_each(function()
+        stub(spawn, "python3", mockx.returns(Result.success()))
+        spawn.python3.on_call_with({ "--version" }).returns(Result.success { stdout = "Python 3.12.0" })
+    end)
+
     it("should init venv without upgrading pip", function()
         local ctx = create_dummy_context()
         stub(ctx, "promote_cwd")