Use already-installed tools in uv tool run (#4750)

## Summary

This doesn't cache the tool environment; rather, it just uses the `tool
install` environment if it satisfies the request.

Closes https://github.com/astral-sh/uv/issues/4742.
This commit is contained in:
Charlie Marsh 2024-07-03 12:35:11 -04:00 committed by GitHub
parent a604f15028
commit f980e3f4fc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 447 additions and 128 deletions

1
Cargo.lock generated
View file

@ -4461,6 +4461,7 @@ dependencies = [
"rustc-hash 2.0.0",
"serde",
"serde_json",
"tempfile",
"textwrap",
"thiserror",
"tikv-jemallocator",

View file

@ -58,6 +58,7 @@ regex = { workspace = true }
rustc-hash = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tempfile = { workspace = true }
textwrap = { workspace = true }
thiserror = { workspace = true }
tokio = { workspace = true }

View file

@ -0,0 +1,48 @@
use pypi_types::Requirement;
use uv_cache::Cache;
use uv_client::Connectivity;
use uv_configuration::{Concurrency, PreviewMode};
use uv_python::Interpreter;
use uv_requirements::RequirementsSpecification;
use crate::commands::{project, SharedState};
use crate::printer::Printer;
use crate::settings::ResolverInstallerSettings;
/// Resolve any [`UnnamedRequirements`].
pub(super) async fn resolve_requirements(
requirements: impl Iterator<Item = &str>,
interpreter: &Interpreter,
settings: &ResolverInstallerSettings,
state: &SharedState,
preview: PreviewMode,
connectivity: Connectivity,
concurrency: Concurrency,
native_tls: bool,
cache: &Cache,
printer: Printer,
) -> anyhow::Result<Vec<Requirement>> {
// Parse the requirements.
let requirements = {
let mut parsed = vec![];
for requirement in requirements {
parsed.push(RequirementsSpecification::parse_package(requirement)?);
}
parsed
};
// Resolve the parsed requirements.
project::resolve_names(
requirements,
interpreter,
settings,
state,
preview,
connectivity,
concurrency,
native_tls,
cache,
printer,
)
.await
}

View file

@ -18,8 +18,7 @@ use uv_fs::Simplified;
use uv_installer::SitePackages;
use uv_normalize::PackageName;
use uv_python::{
EnvironmentPreference, Interpreter, PythonFetch, PythonInstallation, PythonPreference,
PythonRequest,
EnvironmentPreference, PythonFetch, PythonInstallation, PythonPreference, PythonRequest,
};
use uv_requirements::RequirementsSpecification;
use uv_tool::{entrypoint_paths, find_executable_directory, InstalledTools, Tool, ToolEntrypoint};
@ -27,7 +26,8 @@ use uv_warnings::warn_user_once;
use crate::commands::pip::operations::Modifications;
use crate::commands::project::update_environment;
use crate::commands::{project, ExitStatus, SharedState};
use crate::commands::tool::common::resolve_requirements;
use crate::commands::{ExitStatus, SharedState};
use crate::printer::Printer;
use crate::settings::ResolverInstallerSettings;
@ -333,41 +333,3 @@ pub(crate) async fn install(
Ok(ExitStatus::Success)
}
/// Resolve any [`UnnamedRequirements`].
async fn resolve_requirements(
requirements: impl Iterator<Item = &str>,
interpreter: &Interpreter,
settings: &ResolverInstallerSettings,
state: &SharedState,
preview: PreviewMode,
connectivity: Connectivity,
concurrency: Concurrency,
native_tls: bool,
cache: &Cache,
printer: Printer,
) -> Result<Vec<Requirement>> {
// Parse the requirements.
let requirements = {
let mut parsed = vec![];
for requirement in requirements {
parsed.push(RequirementsSpecification::parse_package(requirement)?);
}
parsed
};
// Resolve the parsed requirements.
project::resolve_names(
requirements,
interpreter,
settings,
state,
preview,
connectivity,
concurrency,
native_tls,
cache,
printer,
)
.await
}

View file

@ -1,3 +1,4 @@
mod common;
pub(crate) mod dir;
pub(crate) mod install;
pub(crate) mod list;

View file

@ -1,28 +1,33 @@
use std::borrow::Cow;
use std::ffi::OsString;
use std::ops::Deref;
use std::path::PathBuf;
use std::str::FromStr;
use anyhow::{Context, Result};
use itertools::Itertools;
use pep440_rs::Version;
use tokio::process::Command;
use tracing::debug;
use distribution_types::UnresolvedRequirementSpecification;
use pep440_rs::Version;
use uv_cache::Cache;
use uv_cli::ExternalCommand;
use uv_client::{BaseClientBuilder, Connectivity};
use uv_configuration::{Concurrency, PreviewMode};
use uv_installer::{SatisfiesResult, SitePackages};
use uv_normalize::PackageName;
use uv_python::{
EnvironmentPreference, PythonEnvironment, PythonFetch, PythonInstallation, PythonPreference,
PythonRequest,
};
use uv_requirements::{RequirementsSource, RequirementsSpecification};
use uv_requirements::RequirementsSpecification;
use uv_tool::InstalledTools;
use uv_warnings::warn_user_once;
use crate::commands::pip::operations::Modifications;
use crate::commands::project::update_environment;
use crate::commands::tool::common::resolve_requirements;
use crate::commands::{ExitStatus, SharedState};
use crate::printer::Printer;
use crate::settings::ResolverInstallerSettings;
@ -34,7 +39,7 @@ pub(crate) async fn run(
with: Vec<String>,
python: Option<String>,
settings: ResolverInstallerSettings,
_isolated: bool,
isolated: bool,
preview: PreviewMode,
python_preference: PythonPreference,
python_fetch: PythonFetch,
@ -59,63 +64,23 @@ pub(crate) async fn run(
parse_target(target)?
};
let requirements = [RequirementsSource::from_package(from.to_string())]
.into_iter()
.chain(with.into_iter().map(RequirementsSource::from_package))
.collect::<Vec<_>>();
let client_builder = BaseClientBuilder::new()
.connectivity(connectivity)
.native_tls(native_tls);
let spec =
RequirementsSpecification::from_simple_sources(&requirements, &client_builder).await?;
// TODO(zanieb): When implementing project-level tools, discover the project and check if it has the tool.
// TODO(zanieb): Determine if we should layer on top of the project environment if it is present.
// If necessary, create an environment for the ephemeral requirements.
debug!("Syncing ephemeral environment.");
// Discover an interpreter.
let interpreter = PythonInstallation::find_or_fetch(
python.as_deref().map(PythonRequest::parse),
EnvironmentPreference::OnlySystem,
// Get or create a compatible environment in which to execute the tool.
let environment = get_or_create_environment(
&from,
&with,
python.as_deref(),
&settings,
isolated,
preview,
python_preference,
python_fetch,
&client_builder,
cache,
)
.await?
.into_interpreter();
// Create a virtual environment.
let temp_dir = cache.environment()?;
let venv = uv_virtualenv::create_venv(
temp_dir.path(),
interpreter,
uv_virtualenv::Prompt::None,
false,
false,
)?;
// Install the ephemeral requirements.
let ephemeral_env = Some(
update_environment(
venv,
spec,
Modifications::Sufficient,
&settings,
&SharedState::default(),
preview,
connectivity,
concurrency,
native_tls,
cache,
printer,
)
.await?,
);
.await?;
// TODO(zanieb): Determine the command via the package entry points
let command = target;
@ -126,12 +91,7 @@ pub(crate) async fn run(
// Construct the `PATH` environment variable.
let new_path = std::env::join_paths(
ephemeral_env
.as_ref()
.map(PythonEnvironment::scripts)
.into_iter()
.map(PathBuf::from)
.chain(
std::iter::once(environment.scripts().to_path_buf()).chain(
std::env::var_os("PATH")
.as_ref()
.iter()
@ -142,13 +102,7 @@ pub(crate) async fn run(
// Construct the `PYTHONPATH` environment variable.
let new_python_path = std::env::join_paths(
ephemeral_env
.as_ref()
.map(PythonEnvironment::site_packages)
.into_iter()
.flatten()
.map(PathBuf::from)
.chain(
environment.site_packages().map(PathBuf::from).chain(
std::env::var_os("PYTHONPATH")
.as_ref()
.iter()
@ -180,6 +134,173 @@ pub(crate) async fn run(
}
}
#[derive(Debug)]
enum ToolEnvironment {
Existing(PythonEnvironment),
Ephemeral(PythonEnvironment, #[allow(dead_code)] tempfile::TempDir),
}
impl Deref for ToolEnvironment {
type Target = PythonEnvironment;
fn deref(&self) -> &Self::Target {
match self {
ToolEnvironment::Existing(environment) => environment,
ToolEnvironment::Ephemeral(environment, _) => environment,
}
}
}
/// Get or create a [`PythonEnvironment`] in which to run the specified tools.
///
/// If the target tool is already installed in a compatible environment, returns that
/// [`PythonEnvironment`]. Otherwise, creates an ephemeral environment.
async fn get_or_create_environment(
from: &str,
with: &[String],
python: Option<&str>,
settings: &ResolverInstallerSettings,
isolated: bool,
preview: PreviewMode,
python_preference: PythonPreference,
python_fetch: PythonFetch,
connectivity: Connectivity,
concurrency: Concurrency,
native_tls: bool,
cache: &Cache,
printer: Printer,
) -> Result<ToolEnvironment> {
let client_builder = BaseClientBuilder::new()
.connectivity(connectivity)
.native_tls(native_tls);
let python_request = python.map(PythonRequest::parse);
// Discover an interpreter.
let interpreter = PythonInstallation::find_or_fetch(
python_request.clone(),
EnvironmentPreference::OnlySystem,
python_preference,
python_fetch,
&client_builder,
cache,
)
.await?
.into_interpreter();
// Initialize any shared state.
let state = SharedState::default();
// Resolve the `from` requirement.
let from = {
resolve_requirements(
std::iter::once(from),
&interpreter,
settings,
&state,
preview,
connectivity,
concurrency,
native_tls,
cache,
printer,
)
.await?
.pop()
.unwrap()
};
// Combine the `from` and `with` requirements.
let requirements = {
let mut requirements = Vec::with_capacity(1 + with.len());
requirements.push(from.clone());
requirements.extend(
resolve_requirements(
with.iter().map(String::as_str),
&interpreter,
settings,
&state,
preview,
connectivity,
concurrency,
native_tls,
cache,
printer,
)
.await?,
);
requirements
};
if !isolated {
let installed_tools = InstalledTools::from_settings()?;
// Check if the tool is already installed in a compatible environment.
let existing_environment =
installed_tools
.get_environment(&from.name, cache)?
.filter(|environment| {
python_request.as_ref().map_or(true, |python_request| {
python_request.satisfied(environment.interpreter(), cache)
})
});
if let Some(environment) = existing_environment {
// Check if the installed packages meet the requirements.
let site_packages = SitePackages::from_environment(&environment)?;
let requirements = requirements
.iter()
.cloned()
.map(UnresolvedRequirementSpecification::from)
.collect::<Vec<_>>();
let constraints = [];
if matches!(
site_packages.satisfies(&requirements, &constraints),
Ok(SatisfiesResult::Fresh { .. })
) {
debug!("Using existing tool `{}`", from.name);
return Ok(ToolEnvironment::Existing(environment));
}
}
}
// TODO(zanieb): When implementing project-level tools, discover the project and check if it has the tool.
// TODO(zanieb): Determine if we should layer on top of the project environment if it is present.
// If necessary, create an environment for the ephemeral requirements.
debug!("Syncing ephemeral environment.");
// Create a virtual environment.
let temp_dir = cache.environment()?;
let venv = uv_virtualenv::create_venv(
temp_dir.path(),
interpreter,
uv_virtualenv::Prompt::None,
false,
false,
)?;
// Install the ephemeral requirements.
let spec = RequirementsSpecification::from_requirements(requirements.clone());
let ephemeral_env = update_environment(
venv,
spec,
Modifications::Exact,
settings,
&state,
preview,
connectivity,
concurrency,
native_tls,
cache,
printer,
)
.await?;
Ok(ToolEnvironment::Ephemeral(ephemeral_env, temp_dir))
}
/// Parse a target into a command name and a requirement.
fn parse_target(target: &OsString) -> Result<(Cow<OsString>, Cow<str>)> {
let Some(target_str) = target.to_str() else {

View file

@ -1,5 +1,8 @@
#![cfg(all(feature = "python", feature = "pypi"))]
use assert_cmd::prelude::*;
use assert_fs::prelude::*;
use common::{uv_snapshot, TestContext};
mod common;
@ -7,9 +10,15 @@ mod common;
#[test]
fn tool_run_args() {
let context = TestContext::new("3.12");
let tool_dir = context.temp_dir.child("tools");
let bin_dir = context.temp_dir.child("bin");
// We treat arguments before the command as uv arguments
uv_snapshot!(context.filters(), context.tool_run().arg("--version").arg("pytest"), @r###"
uv_snapshot!(context.filters(), context.tool_run()
.arg("--version")
.arg("pytest")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: true
exit_code: 0
----- stdout -----
@ -19,7 +28,11 @@ fn tool_run_args() {
"###);
// We don't treat arguments after the command as uv arguments
uv_snapshot!(context.filters(), context.tool_run().arg("pytest").arg("--version"), @r###"
uv_snapshot!(context.filters(), context.tool_run()
.arg("pytest")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: true
exit_code: 0
----- stdout -----
@ -37,7 +50,12 @@ fn tool_run_args() {
"###);
// Can use `--` to separate uv arguments from the command arguments.
uv_snapshot!(context.filters(), context.tool_run().arg("--").arg("pytest").arg("--version"), @r###"
uv_snapshot!(context.filters(), context.tool_run()
.arg("--")
.arg("pytest")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: true
exit_code: 0
----- stdout -----
@ -57,8 +75,14 @@ fn tool_run_args() {
#[test]
fn tool_run_at_version() {
let context = TestContext::new("3.12");
let tool_dir = context.temp_dir.child("tools");
let bin_dir = context.temp_dir.child("bin");
uv_snapshot!(context.filters(), context.tool_run().arg("pytest@8.0.0").arg("--version"), @r###"
uv_snapshot!(context.filters(), context.tool_run()
.arg("pytest@8.0.0")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: true
exit_code: 0
----- stdout -----
@ -76,7 +100,11 @@ fn tool_run_at_version() {
"###);
// Empty versions are just treated as package and command names
uv_snapshot!(context.filters(), context.tool_run().arg("pytest@").arg("--version"), @r###"
uv_snapshot!(context.filters(), context.tool_run()
.arg("pytest@")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: false
exit_code: 2
----- stdout -----
@ -90,7 +118,11 @@ fn tool_run_at_version() {
"###);
// Invalid versions are just treated as package and command names
uv_snapshot!(context.filters(), context.tool_run().arg("pytest@invalid").arg("--version"), @r###"
uv_snapshot!(context.filters(), context.tool_run()
.arg("pytest@invalid")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: false
exit_code: 2
----- stdout -----
@ -111,7 +143,13 @@ fn tool_run_at_version() {
.collect::<Vec<_>>();
// When `--from` is used, `@` is not treated as a version request
uv_snapshot!(filters, context.tool_run().arg("--from").arg("pytest").arg("pytest@8.0.0").arg("--version"), @r###"
uv_snapshot!(filters, context.tool_run()
.arg("--from")
.arg("pytest")
.arg("pytest@8.0.0")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: false
exit_code: 2
----- stdout -----
@ -133,8 +171,16 @@ fn tool_run_at_version() {
#[test]
fn tool_run_from_version() {
let context = TestContext::new("3.12");
let tool_dir = context.temp_dir.child("tools");
let bin_dir = context.temp_dir.child("bin");
uv_snapshot!(context.filters(), context.tool_run().arg("--from").arg("pytest==8.0.0").arg("pytest").arg("--version"), @r###"
uv_snapshot!(context.filters(), context.tool_run()
.arg("--from")
.arg("pytest==8.0.0")
.arg("pytest")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: true
exit_code: 0
----- stdout -----
@ -151,3 +197,142 @@ fn tool_run_from_version() {
+ pytest==8.0.0
"###);
}
#[test]
fn tool_run_from_install() {
let context = TestContext::new("3.12");
let tool_dir = context.temp_dir.child("tools");
let bin_dir = context.temp_dir.child("bin");
// Install `black` at a specific version.
context
.tool_install()
.arg("black==24.1.0")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str())
.assert()
.success();
// Verify that `tool run black` uses the already-installed version.
uv_snapshot!(context.filters(), context.tool_run()
.arg("black")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: true
exit_code: 0
----- stdout -----
black, 24.1.0 (compiled: yes)
Python (CPython) 3.12.[X]
----- stderr -----
warning: `uv tool run` is experimental and may change without warning.
"###);
// Verify that `--isolated` uses an isolated environment.
uv_snapshot!(context.filters(), context.tool_run()
.arg("--isolated")
.arg("black")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: true
exit_code: 0
----- stdout -----
black, 24.3.0 (compiled: yes)
Python (CPython) 3.12.[X]
----- stderr -----
warning: `uv tool run` is experimental and may change without warning.
Resolved 6 packages in [TIME]
Prepared 1 package in [TIME]
Installed 6 packages in [TIME]
+ black==24.3.0
+ click==8.1.7
+ mypy-extensions==1.0.0
+ packaging==24.0
+ pathspec==0.12.1
+ platformdirs==4.2.0
"###);
// Verify that `tool run black` at a different version installs the new version.
uv_snapshot!(context.filters(), context.tool_run()
.arg("black@24.1.1")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: true
exit_code: 0
----- stdout -----
black, 24.1.1 (compiled: yes)
Python (CPython) 3.12.[X]
----- stderr -----
warning: `uv tool run` is experimental and may change without warning.
Resolved 6 packages in [TIME]
Prepared 1 package in [TIME]
Installed 6 packages in [TIME]
+ black==24.1.1
+ click==8.1.7
+ mypy-extensions==1.0.0
+ packaging==24.0
+ pathspec==0.12.1
+ platformdirs==4.2.0
"###);
// Verify that `tool run black` at a different version (via `--from`) installs the new version.
uv_snapshot!(context.filters(), context.tool_run()
.arg("--from")
.arg("black==24.1.1")
.arg("black")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: true
exit_code: 0
----- stdout -----
black, 24.1.1 (compiled: yes)
Python (CPython) 3.12.[X]
----- stderr -----
warning: `uv tool run` is experimental and may change without warning.
Resolved 6 packages in [TIME]
Installed 6 packages in [TIME]
+ black==24.1.1
+ click==8.1.7
+ mypy-extensions==1.0.0
+ packaging==24.0
+ pathspec==0.12.1
+ platformdirs==4.2.0
"###);
// Verify that `--with` installs a new version.
// TODO(charlie): This could (in theory) layer the `--with` requirements on top of the existing
// environment.
uv_snapshot!(context.filters(), context.tool_run()
.arg("--with")
.arg("iniconfig")
.arg("black")
.arg("--version")
.env("UV_TOOL_DIR", tool_dir.as_os_str())
.env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
success: true
exit_code: 0
----- stdout -----
black, 24.3.0 (compiled: yes)
Python (CPython) 3.12.[X]
----- stderr -----
warning: `uv tool run` is experimental and may change without warning.
Resolved 7 packages in [TIME]
Prepared 1 package in [TIME]
Installed 7 packages in [TIME]
+ black==24.3.0
+ click==8.1.7
+ iniconfig==2.0.0
+ mypy-extensions==1.0.0
+ packaging==24.0
+ pathspec==0.12.1
+ platformdirs==4.2.0
"###);
}