Skip to content

Commit

Permalink
feat: add uv as installer for python
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
KingMichaelPark committed Sep 13, 2024
1 parent e2f7f90 commit b1bceba
Show file tree
Hide file tree
Showing 6 changed files with 59 additions and 13 deletions.
4 changes: 4 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
4 changes: 4 additions & 0 deletions doc/mason.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
56 changes: 43 additions & 13 deletions lua/mason-core/installer/managers/pypi.lua
Original file line number Diff line number Diff line change
Expand Up @@ -9,22 +9,32 @@ 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 use_uv = settings.current.pip.use_uv
local VENV_DIR = "venv"

---@async
---@param candidates string[]
local function resolve_python3(candidates)
local is_executable = _.compose(_.equals(1), vim.fn.executable)
a.scheduler()
if use_uv then
candidates = { "uv" }
end
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+)")
local ok, version
if use_uv then
ok, version = pcall(semver.new, version_output:match "uv (%d+.%d+.%d+)")
else
ok, version = pcall(semver.new, version_output:match "Python (3%.%d+.%d+)")
end
if ok then
return { executable = candidate, version = version }
end
Expand Down Expand Up @@ -127,7 +137,14 @@ local function create_venv(pkg)

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
Expand Down Expand Up @@ -162,16 +179,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,
}
-- vim.api.nvim_set_current_dir(curdir)
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
Expand All @@ -185,7 +215,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
Expand Down
3 changes: 3 additions & 0 deletions lua/mason-core/installer/registry/providers/pypi.lua
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
}

Expand All @@ -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
Expand Down
4 changes: 4 additions & 0 deletions lua/mason/settings.lua
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down

0 comments on commit b1bceba

Please sign in to comment.