feat: add uv as installer for python

UV is a very fast installer for python packages that can be 10-100x faster to resolve packages. This adds an option for Mason to use it instead of pip to resolve python packages that are installed via Mason. More info about the replacement: https://github.com/astral-sh/uv I have no relationship with uv, it is just very fast and it would be nice to have updates for packages like sqlfluff take a lot less time than they currently do to resolve during updates. fix: ensure the virtual environment is .venv for uv
williamboman · Oct 25, 2024 · db211a4 · db211a4
1 parent e2f7f90
commit db211a4
Show file tree

Hide file tree

Showing 6 changed files with 76 additions and 22 deletions.
diff --git a/README.md b/README.md
@@ -252,6 +252,10 @@ local DEFAULT_SETTINGS = {
         -- Whether to upgrade pip to the latest version in the virtual environment before installing packages.
         upgrade_pip = false,
 
+        ---@since 1.8.0
+        -- Whether to use uv to install packages instead of pip
+        use_uv      = false,
+
         ---@since 1.0.0
         -- These args will be added to `pip install` calls. Note that setting extra args might impact intended behavior
         -- and is not recommended.

diff --git a/doc/mason.txt b/doc/mason.txt
@@ -314,6 +314,10 @@ Example:
             -- Whether to upgrade pip to the latest version in the virtual environment before installing packages.
             upgrade_pip = false,
 
+            ---@since 1.8.0
+            -- Whether to use uv to install packages instead of pip
+            use_uv      = false,
+
             ---@since 1.0.0
             -- These args will be added to `pip install` calls. Note that setting extra args might impact intended behavior
             -- and is not recommended.

diff --git a/lua/mason-core/installer/managers/pypi.lua b/lua/mason-core/installer/managers/pypi.lua
@@ -9,11 +9,18 @@ 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 settings = require "mason.settings"
 local spawn = require "mason-core.spawn"
 
 local M = {}
 
-local VENV_DIR = "venv"
+local use_uv = settings.current.pip.use_uv
+local VENV_DIR
+if use_uv then
+    VENV_DIR = ".venv"
+else
+    VENV_DIR = "venv"
+end
 
 ---@async
 ---@param candidates string[]
@@ -22,11 +29,20 @@ 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 }
+        if use_uv and candidate == "uv" then
+            ---@type string
+            local version_output = spawn[candidate]({ "--version" }):map(_.prop "stdout"):get_or_else ""
+            local ok, version = pcall(semver.new, version_output:match "uv (%d+.%d+.%d+).*")
+            if ok then
+                return { executable = candidate, version = version }
+            end
+        elseif not use_uv then
+            ---@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
     end
     return nil
@@ -76,14 +92,14 @@ local function create_venv(pkg)
     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_candidates = platform.is.win and { "python", "python3", "uv" } or { "python3", "python", "uv" }
     local stock_target = resolve_python3(stock_candidates)
     if stock_target then
         log.fmt_debug("Resolved stock python3 installation version %s", stock_target.version)
     end
 
     -- 2. Resolve suitable versioned python3 installation (python3.12, python3.11, etc.).
-    local versioned_candidates = {}
+    local versioned_candidates = { "uv" }
     if supported_python_versions ~= nil then
         if stock_target and not pep440_check_version(tostring(stock_target.version), supported_python_versions) then
             log.fmt_debug("Finding versioned candidates for %s", supported_python_versions)
@@ -103,7 +119,8 @@ local function create_venv(pkg)
     -- 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
+        use_uv == false
+        and target == stock_target
         and supported_python_versions ~= nil
         and not pep440_check_version(tostring(target.version), supported_python_versions)
     then
@@ -125,9 +142,14 @@ local function create_venv(pkg)
         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", "--system-site-packages", VENV_DIR }
+    if use_uv then
+        log.fmt_debug("Found uv installation version=%s, executable=%s", target.version, target.executable)
+        return ctx.spawn[target.executable] { "venv", VENV_DIR }
+    else
+        log.fmt_debug("Found python3 installation version=%s, executable=%s", target.version, target.executable)
+        return ctx.spawn[target.executable] { "-m", "venv", "--system-site-packages", VENV_DIR }
+    end
 end
 
 ---@param ctx InstallContext
@@ -153,6 +175,9 @@ end
 ---@param args SpawnArgs
 local function venv_python(args)
     local ctx = installer.context()
+    if use_uv then
+        return ctx.spawn[{ "uv", "venv" }](args)
+    end
     return find_venv_executable(ctx, "python"):and_then(function(python_path)
         return ctx.spawn[path.concat { ctx.cwd:get(), python_path }](args)
     end)
@@ -162,16 +187,29 @@ end
 ---@param pkgs string[]
 ---@param extra_args? string[]
 local function pip_install(pkgs, extra_args)
-    return venv_python {
-        "-m",
-        "pip",
-        "--disable-pip-version-check",
-        "install",
-        "--ignore-installed",
-        "-U",
-        extra_args or vim.NIL,
-        pkgs,
-    }
+    if use_uv then
+        local ctx = installer.context()
+
+        local task = ctx.spawn["uv"] {
+            "pip",
+            "install",
+            "-U",
+            extra_args or vim.NIL,
+            pkgs,
+        }
+        return task
+    else
+        return venv_python {
+            "-m",
+            "pip",
+            "--disable-pip-version-check",
+            "install",
+            "--ignore-installed",
+            "-U",
+            extra_args or vim.NIL,
+            pkgs,
+        }
+    end
 end
 
 ---@async
@@ -185,7 +223,7 @@ function M.init(opts)
         ctx:promote_cwd()
         try(create_venv(opts.package))
 
-        if opts.upgrade_pip then
+        if opts.upgrade_pip and not use_uv then
             ctx.stdio_sink.stdout "Upgrading pip inside the virtual environment…\n"
             try(pip_install({ "pip" }, opts.install_extra_args))
         end

diff --git a/lua/mason-core/installer/registry/providers/pypi.lua b/lua/mason-core/installer/registry/providers/pypi.lua
@@ -27,6 +27,7 @@ function M.parse(source, purl)
             pip = {
                 upgrade = settings.current.pip.upgrade_pip,
                 extra_args = settings.current.pip.install_args,
+                use_uv = settings.current.pip.use_uv,
             },
         }
 
@@ -48,11 +49,13 @@ function M.install(ctx, source)
             },
             upgrade_pip = source.pip.upgrade,
             install_extra_args = source.pip.extra_args,
+            use_uv = source.pip.use_uv,
         })
         try(pypi.install(source.package, source.version, {
             extra = source.extra,
             extra_packages = source.extra_packages,
             install_extra_args = source.pip.extra_args,
+            use_uv = source.pip.use_uv,
         }))
     end)
 end

diff --git a/lua/mason/settings.lua b/lua/mason/settings.lua
@@ -60,6 +60,10 @@ local DEFAULT_SETTINGS = {
         -- Whether to upgrade pip to the latest version in the virtual environment before installing packages.
         upgrade_pip = false,
 
+        ---@since 1.8.0
+        -- Whether to use uv to install packages instead of pip
+        use_uv = false,
+
         ---@since 1.0.0
         -- These args will be added to `pip install` calls. Note that setting extra args might impact intended behavior
         -- and is not recommended.

diff --git a/tests/mason-core/installer/registry/providers/pypi_spec.lua b/tests/mason-core/installer/registry/providers/pypi_spec.lua
@@ -31,6 +31,7 @@ describe("pypi provider :: parsing", function()
                 pip = {
                     upgrade = true,
                     extra_args = { "--proxy", "http://localghost" },
+                    use_uv = false,
                 },
             },
             pypi.parse({ extra_packages = { "extra" } }, purl())