feat(pypi): improve resolving suitable python version

williamboman · williamboman · commit 588130b919c5 · 2024-05-31T14:29:15.000+02:00
diff --git a/lua/mason-core/installer/managers/pypi.lua b/lua/mason-core/installer/managers/pypi.lua
@@ -5,19 +5,20 @@ 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"
 
 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 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,16 +69,27 @@ 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(
@@ -69,6 +98,22 @@ local function create_venv()
         )
         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
+        ctx.stdio_sink.stderr(
+            ("Warning: The resolved Python version %s is not compatible with the required Python versions: %s.\n"):format(
+                target.version,
+                supported_python_versions
+            )
+        )
+    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,15 +163,15 @@ 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)
         local ctx = installer.context()
 
         -- 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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" 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
@@ -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" 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)
diff --git a/tests/mason-core/installer/managers/pypi_spec.lua b/tests/mason-core/installer/managers/pypi_spec.lua
@@ -2,6 +2,7 @@ local Result = require "mason-core.result"
 local installer = require "mason-core.installer"
 local match = require "luassert.match"
 local path = require "mason-core.path"
+local providers = require "mason-core.providers"
 local pypi = require "mason-core.installer.managers.pypi"
 local spawn = require "mason-core.spawn"
 local spy = require "luassert.spy"
@@ -26,9 +27,10 @@ describe("pypi manager", function()
     it("should init venv without upgrading pip", function()
         local ctx = create_dummy_context()
         stub(ctx, "promote_cwd")
+        stub(providers.pypi, "get_supported_python_versions", mockx.returns(Result.failure()))
 
         installer.exec_in_context(ctx, function()
-            pypi.init { upgrade_pip = false }
+            pypi.init { package = { name = "cmake-language-server", version = "0.1.10" }, upgrade_pip = false }
         end)
 
         assert.spy(ctx.promote_cwd).was_called(1)
@@ -44,10 +46,15 @@ describe("pypi manager", function()
         local ctx = create_dummy_context()
         stub(ctx, "promote_cwd")
         stub(ctx.fs, "file_exists")
+        stub(providers.pypi, "get_supported_python_versions", mockx.returns(Result.failure()))
         ctx.fs.file_exists.on_call_with(match.ref(ctx.fs), "venv/bin/python").returns(true)
 
         installer.exec_in_context(ctx, function()
-            pypi.init { upgrade_pip = true, install_extra_args = { "--proxy", "http://localhost" } }
+            pypi.init {
+                package = { name = "cmake-language-server", version = "0.1.10" },
+                upgrade_pip = true,
+                install_extra_args = { "--proxy", "http://localhost" },
+            }
         end)
 
         assert.spy(ctx.promote_cwd).was_called(1)
@@ -69,6 +76,67 @@ describe("pypi manager", function()
         }
     end)
 
+    it("should find versioned candidates during init", function()
+        local ctx = create_dummy_context()
+        stub(ctx, "promote_cwd")
+        stub(ctx.fs, "file_exists")
+        stub(providers.pypi, "get_supported_python_versions", mockx.returns(Result.success ">=3.12"))
+        stub(spawn, "python3.12")
+        spawn["python3.12"].on_call_with({ "--version" }).returns(Result.success { stdout = "Python 3.12.0" })
+        ctx.fs.file_exists.on_call_with(match.ref(ctx.fs), "venv/bin/python").returns(true)
+
+        installer.exec_in_context(ctx, function()
+            pypi.init {
+                package = { name = "cmake-language-server", version = "0.1.10" },
+                upgrade_pip = true,
+                install_extra_args = { "--proxy", "http://localhost" },
+            }
+        end)
+
+        assert.spy(ctx.promote_cwd).was_called(1)
+        assert.spy(ctx.spawn["python3.12"]).was_called(1)
+        assert.spy(ctx.spawn["python3.12"]).was_called_with {
+            "-m",
+            "venv",
+            "venv",
+        }
+    end)
+
+    it("should default to stock version if unable to find suitable versioned candidate during init", function()
+        local ctx = create_dummy_context()
+        spy.on(ctx.stdio_sink, "stderr")
+        stub(ctx, "promote_cwd")
+        stub(ctx.fs, "file_exists")
+        stub(providers.pypi, "get_supported_python_versions", mockx.returns(Result.success ">=3.8"))
+        stub(vim.fn, "executable")
+        vim.fn.executable.on_call_with("python3.12").returns(0)
+        vim.fn.executable.on_call_with("python3.11").returns(0)
+        vim.fn.executable.on_call_with("python3.10").returns(0)
+        vim.fn.executable.on_call_with("python3.9").returns(0)
+        vim.fn.executable.on_call_with("python3.8").returns(0)
+        stub(spawn, "python3", mockx.returns(Result.success()))
+        spawn.python3.on_call_with({ "--version" }).returns(Result.success { stdout = "Python 3.5.0" })
+
+        installer.exec_in_context(ctx, function()
+            pypi.init {
+                package = { name = "cmake-language-server", version = "0.1.10" },
+                upgrade_pip = true,
+                install_extra_args = { "--proxy", "http://localhost" },
+            }
+        end)
+
+        assert.spy(ctx.promote_cwd).was_called(1)
+        assert.spy(ctx.spawn.python3).was_called(1)
+        assert.spy(ctx.spawn.python3).was_called_with {
+            "-m",
+            "venv",
+            "venv",
+        }
+        assert
+            .spy(ctx.stdio_sink.stderr)
+            .was_called_with "Warning: The resolved Python version 3.5.0 is not compatible with the required Python versions: >=3.8.\n"
+    end)
+
     it("should install", function()
         local ctx = create_dummy_context()
         stub(ctx.fs, "file_exists")
diff --git a/tests/mason-core/installer/registry/providers/pypi_spec.lua b/tests/mason-core/installer/registry/providers/pypi_spec.lua

Original file line number	Diff line number	Diff line change
`@@ -83,6 +83,8 @@ api.pypi = {`
`83`	`83`	`latest = get "/api/pypi/{package}/versions/latest",`
`84`	`84`	`---@type ApiSignature<{ package: string }>`
`85`	`85`	`all = get "/api/pypi/{package}/versions/all",`
	`86`	`+ ---@type ApiSignature<{ package: string, version: string }>`
	`87`	`+ get = get "/api/pypi/{package}/versions/{version}",`
`86`	`88`	`},`
`87`	`89`	`}`
`88`	`90`