feat(pypi): improve resolving suitable python version (#1725)

7 files changed, 170 insertions, 17 deletions

diff --git a/lua/mason-core/installer/managers/pypi.lua b/lua/mason-core/installer/managers/pypi.lua
index 9cd40ea3..e9d54559 100644
--- a/lua/mason-core/installer/managers/pypi.lua
+++ b/lua/mason-core/installer/managers/pypi.lua
@@ -5,7 +5,9 @@ local a = require "mason-core.async"
 local installer = require "mason-core.installer"
 local log = require "mason-core.log"
 local path = require "mason-core.path"
+local pep440 = require "mason-core.pep440"
 local platform = require "mason-core.platform"
+local providers = require "mason-core.providers"
 local semver = require "mason-core.semver"
 local spawn = require "mason-core.spawn"
 
@@ -13,11 +15,10 @@ local M = {}
 
 local VENV_DIR = "venv"
 
-local is_executable = _.compose(_.equals(1), vim.fn.executable)
-
 ---@async
 ---@param candidates string[]
 local function resolve_python3(candidates)
+    local is_executable = _.compose(_.equals(1), vim.fn.executable)
     a.scheduler()
     local available_candidates = _.filter(is_executable, candidates)
     for __, candidate in ipairs(available_candidates) do
@@ -31,16 +32,33 @@ local function resolve_python3(candidates)
     return nil
 end
 
----@param min_version? Semver
-local function get_versioned_candidates(min_version)
+---@param version string
+---@param specifiers string
+local function pep440_check_version(version, specifiers)
+    -- The version check only implements a subset of the PEP440 specification and may error with certain inputs.
+    local ok, result = pcall(pep440.check_version, version, specifiers)
+    if not ok then
+        log.fmt_warn(
+            "Failed to check PEP440 version compatibility for version %s with specifiers %s: %s",
+            version,
+            specifiers,
+            result
+        )
+        return false
+    end
+    return result
+end
+
+---@param supported_python_versions string
+local function get_versioned_candidates(supported_python_versions)
     return _.filter_map(function(pair)
         local version, executable = unpack(pair)
-        if not min_version or version > min_version then
-            return Optional.of(executable)
-        else
+        if not pep440_check_version(tostring(version), supported_python_versions) then
             return Optional.empty()
         end
+        return Optional.of(executable)
     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" },
@@ -51,24 +69,60 @@ local function get_versioned_candidates(min_version)
 end
 
 ---@async
-local function create_venv()
+---@param pkg { name: string, version: string }
+local function create_venv(pkg)
+    local ctx = installer.context()
+    ---@type string?
+    local supported_python_versions = providers.pypi.get_supported_python_versions(pkg.name, pkg.version):get_or_nil()
+
+    -- 1. Resolve stock python3 installation.
     local stock_candidates = platform.is.win and { "python", "python3" } or { "python3", "python" }
     local stock_target = resolve_python3(stock_candidates)
     if stock_target then
         log.fmt_debug("Resolved stock python3 installation version %s", stock_target.version)
     end
-    local versioned_candidates = get_versioned_candidates(stock_target and stock_target.version)
-    log.debug("Resolving versioned python3 candidates", versioned_candidates)
+
+    -- 2. Resolve suitable versioned python3 installation (python3.12, python3.11, etc.).
+    local versioned_candidates = {}
+    if supported_python_versions ~= nil then
+        log.fmt_debug("Finding versioned candidates for %s", supported_python_versions)
+        versioned_candidates = get_versioned_candidates(supported_python_versions)
+    end
     local target = resolve_python3(versioned_candidates) or stock_target
-    local ctx = installer.context()
+
     if not target then
-        ctx.stdio_sink.stderr(
-            ("Unable to find python3 installation. Tried the following candidates: %s.\n"):format(
+        return Result.failure(
+            ("Unable to find python3 installation in PATH. Tried the following candidates: %s."):format(
                 _.join(", ", _.concat(stock_candidates, versioned_candidates))
             )
         )
-        return Result.failure "Failed to find python3 installation."
     end
+
+    -- 3. If a versioned python3 installation was not found, warn the user if the stock python3 installation is outside
+    -- the supported version range.
+    if
+        target == stock_target
+        and supported_python_versions ~= nil
+        and not pep440_check_version(tostring(target.version), supported_python_versions)
+    then
+        if ctx.opts.force then
+            ctx.stdio_sink.stderr(
+                ("Warning: The resolved python3 version %s is not compatible with the required Python versions: %s.\n"):format(
+                    target.version,
+                    supported_python_versions
+                )
+            )
+        else
+            ctx.stdio_sink.stderr "Run with :MasonInstall --force to bypass this version validation.\n"
+            return Result.failure(
+                ("Failed to find a python3 installation in PATH that meets the required versions (%s). Found version: %s."):format(
+                    supported_python_versions,
+                    target.version
+                )
+            )
+        end
+    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 }
@@ -118,7 +172,7 @@ local function pip_install(pkgs, extra_args)
 end
 
 ---@async
----@param opts { upgrade_pip: boolean, install_extra_args?: string[] }
+---@param opts { package: { name: string, version: string }, upgrade_pip: boolean, install_extra_args?: string[] }
 function M.init(opts)
     return Result.try(function(try)
         log.fmt_debug("pypi: init", opts)
@@ -126,7 +180,7 @@ function M.init(opts)
 
         -- 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()
-        try(create_venv())
+        try(create_venv(opts.package))
 
         if opts.upgrade_pip then
             ctx.stdio_sink.stdout "Upgrading pip inside the virtual environment…\n"
diff --git a/lua/mason-core/installer/registry/providers/pypi.lua b/lua/mason-core/installer/registry/providers/pypi.lua
index c162c120..3fe6f89e 100644
--- a/lua/mason-core/installer/registry/providers/pypi.lua
+++ b/lua/mason-core/installer/registry/providers/pypi.lua
@@ -21,7 +21,7 @@ function M.parse(source, purl)
         ---@class ParsedPypiSource : ParsedPackageSource
         local parsed_source = {
             package = purl.name,
-            version = purl.version,
+            version = purl.version --[[ @as string ]],
             extra = _.path({ "qualifiers", "extra" }, purl),
             extra_packages = source.extra_packages,
             pip = {
@@ -42,6 +42,10 @@ function M.install(ctx, source)
 
     return Result.try(function(try)
         try(pypi.init {
+            package = {
+                name = source.package,
+                version = source.version,
+            },
             upgrade_pip = source.pip.upgrade,
             install_extra_args = source.pip.extra_args,
         })
diff --git a/lua/mason-core/pep440/init.lua b/lua/mason-core/pep440/init.lua
new file mode 100644
index 00000000..0141c501
--- /dev/null
+++ b/lua/mason-core/pep440/init.lua
@@ -0,0 +1,64 @@
+-- Function to split a version string into its components
+local function split_version(version)
+    local parts = {}
+    for part in version:gmatch "[^.]+" do
+        table.insert(parts, tonumber(part) or part)
+    end
+    return parts
+end
+
+-- Function to compare two versions
+local function compare_versions(version1, version2)
+    local v1_parts = split_version(version1)
+    local v2_parts = split_version(version2)
+    local len = math.max(#v1_parts, #v2_parts)
+
+    for i = 1, len do
+        local v1_part = v1_parts[i] or 0
+        local v2_part = v2_parts[i] or 0
+
+        if v1_part < v2_part then
+            return -1
+        elseif v1_part > v2_part then
+            return 1
+        end
+    end
+
+    return 0
+end
+
+-- Function to check a version against a single specifier
+local function check_single_specifier(version, specifier)
+    local operator, spec_version = specifier:match "^([<>=!]+)%s*(.+)$"
+    local comp_result = compare_versions(version, spec_version)
+
+    if operator == "==" then
+        return comp_result == 0
+    elseif operator == "!=" then
+        return comp_result ~= 0
+    elseif operator == "<=" then
+        return comp_result <= 0
+    elseif operator == "<" then
+        return comp_result < 0
+    elseif operator == ">=" then
+        return comp_result >= 0
+    elseif operator == ">" then
+        return comp_result > 0
+    else
+        error("Invalid operator in version specifier: " .. operator)
+    end
+end
+
+-- Function to check a version against multiple specifiers
+local function check_version(version, specifiers)
+    for specifier in specifiers:gmatch "[^,]+" do
+        if not check_single_specifier(version, specifier:match "^%s*(.-)%s*$") then
+            return false
+        end
+    end
+    return true
+end
+
+return {
+    check_version = check_version,
+}
diff --git a/lua/mason-core/providers/init.lua b/lua/mason-core/providers/init.lua
index a97d1b60..5e2a8ea0 100644
--- a/lua/mason-core/providers/init.lua
+++ b/lua/mason-core/providers/init.lua
@@ -22,6 +22,7 @@ local settings = require "mason.settings"
 ---@class PyPiProvider
 ---@field get_latest_version? async fun(pkg: string): Result # Result<PyPiPackage>
 ---@field get_all_versions? async fun(pkg: string): Result # Result<string[]> # Sorting should not be relied upon due to "proprietary" sorting algo in pip that is difficult to replicate in mason-registry-api.
+---@field get_supported_python_versions? async fun(pkg: string, version: string): Result # Result<string> # Returns a version specifier as provided by the PyPI API (see PEP440).
 
 ---@alias RubyGem { name: string, version: string }
 
diff --git a/lua/mason-registry/api.lua b/lua/mason-registry/api.lua
index b4acea63..be5d4d92 100644
--- a/lua/mason-registry/api.lua
+++ b/lua/mason-registry/api.lua
@@ -83,6 +83,8 @@ api.pypi = {
         latest = get "/api/pypi/{package}/versions/latest",
         ---@type ApiSignature<{ package: string }>
         all = get "/api/pypi/{package}/versions/all",
+        ---@type ApiSignature<{ package: string, version: string }>
+        get = get "/api/pypi/{package}/versions/{version}",
     },
 }
 
diff --git a/lua/mason/providers/client/pypi.lua b/lua/mason/providers/client/pypi.lua
index ecbbfd56..08e4dbae 100644
--- a/lua/mason/providers/client/pypi.lua
+++ b/lua/mason/providers/client/pypi.lua
@@ -1,6 +1,8 @@
 local Optional = require "mason-core.optional"
+local Result = require "mason-core.result"
 local _ = require "mason-core.functional"
 local a = require "mason-core.async"
+local fetch = require "mason-core.fetch"
 local fs = require "mason-core.fs"
 local platform = require "mason-core.platform"
 local spawn = require "mason-core.spawn"
@@ -50,4 +52,16 @@ return {
         return get_all_versions(pkg):map(_.compose(Optional.of_nilable, _.last)):and_then(synthesize_pkg(pkg))
     end,
     get_all_versions = get_all_versions,
+    get_supported_python_versions = function(pkg, version)
+        return fetch(("https://pypi.org/pypi/%s/%s/json"):format(pkg, version))
+            :map_catching(vim.json.decode)
+            :map(_.path { "info", "requires_python" })
+            :and_then(function(requires_python)
+                if type(requires_python) ~= "string" or requires_python == "" then
+                    return Result.failure "Package does not specify supported Python versions."
+                else
+                    return Result.success(requires_python)
+                end
+            end)
+    end,
 }
diff --git a/lua/mason/providers/registry-api/init.lua b/lua/mason/providers/registry-api/init.lua
index d8802124..26e1c0cb 100644
--- a/lua/mason/providers/registry-api/init.lua
+++ b/lua/mason/providers/registry-api/init.lua
@@ -1,3 +1,5 @@
+local Result = require "mason-core.result"
+local _ = require "mason-core.functional"
 local api = require "mason-registry.api"
 
 ---@type Provider
@@ -31,6 +33,18 @@ return {
         get_all_versions = function(pkg)
             return api.pypi.versions.all { package = pkg }
         end,
+        get_supported_python_versions = function(pkg, version)
+            return api.pypi.versions
+                .get({ package = pkg, version = version })
+                :map(_.prop "requires_python")
+                :and_then(function(requires_python)
+                    if type(requires_python) ~= "string" or requires_python == "" then
+                        return Result.failure "Package does not specify supported Python versions."
+                    else
+                        return Result.success(requires_python)
+                    end
+                end)
+        end,
     },
     rubygems = {
         get_latest_version = function(gem)