Add `uv run --with <pkg>` to run a command with ephemeral requirements (#3077)

Holy cow does installation / resolution take a ton of options. We
side-step most of them here.

If the current environment satisfies the requirements, it is used.
Otherwise, we create a new environment with the requested dependencies.
This commit is contained in:
Zanie Blue 2024-04-19 09:23:26 -05:00 committed by GitHub
parent 7a163ba9f1
commit becb12642a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 574 additions and 33 deletions

View File

@ -1637,6 +1637,10 @@ pub(crate) struct RunArgs {
/// Always use a new virtual environment.
#[arg(long)]
pub(crate) isolated: bool,
/// Run with the given packages installed.
#[arg(long)]
pub(crate) with: Vec<String>,
}
#[derive(Args)]

View File

@ -1,29 +1,55 @@
use std::ffi::OsString;
use std::{env, iter};
use anyhow::Result;
use owo_colors::OwoColorize;
use tempfile::{tempdir_in, TempDir};
use tracing::debug;
use uv_fs::Simplified;
use uv_interpreter::PythonEnvironment;
use crate::commands::reporters::{DownloadReporter, InstallReporter, ResolverReporter};
use crate::commands::ExitStatus;
use crate::commands::{elapsed, ChangeEvent, ChangeEventKind};
use crate::printer::Printer;
use anyhow::{Context, Result};
use distribution_types::{IndexLocations, InstalledMetadata, LocalDist, Name, Resolution};
use install_wheel_rs::linker::LinkMode;
use itertools::Itertools;
use owo_colors::OwoColorize;
use pep508_rs::{MarkerEnvironment, PackageName, Requirement};
use platform_tags::Tags;
use pypi_types::Yanked;
use std::ffi::OsString;
use std::fmt::Write;
use std::{env, iter};
use tempfile::{tempdir_in, TempDir};
use tokio::process::Command;
use tracing::debug;
use uv_cache::Cache;
use uv_client::{BaseClientBuilder, RegistryClient, RegistryClientBuilder};
use uv_configuration::{
ConfigSettings, Constraints, NoBinary, NoBuild, Overrides, Reinstall, SetupPyStrategy,
};
use uv_dispatch::BuildDispatch;
use uv_fs::Simplified;
use uv_installer::{Downloader, Plan, Planner, SitePackages};
use uv_interpreter::{Interpreter, PythonEnvironment};
use uv_requirements::{
ExtrasSpecification, LookaheadResolver, NamedRequirementsResolver, RequirementsSource,
RequirementsSpecification, SourceTreeResolver,
};
use uv_resolver::{
Exclusions, FlatIndex, InMemoryIndex, Manifest, Options, OptionsBuilder, ResolutionGraph,
Resolver,
};
use uv_types::{BuildIsolation, EmptyInstalledPackages, HashStrategy, InFlight};
/// Run a command.
#[allow(clippy::unnecessary_wraps, clippy::too_many_arguments)]
pub(crate) async fn run(
command: String,
args: Vec<String>,
requirements: &[RequirementsSource],
isolated: bool,
cache: &Cache,
printer: Printer,
) -> Result<ExitStatus> {
// Detect the current Python interpreter.
// TODO(zanieb): Create ephemeral environments
// TODO(zanieb): Accept `--python`
let run_env = environment_for_run(isolated, cache)?;
let run_env = environment_for_run(requirements, isolated, cache, printer).await?;
let python_env = run_env.python;
// Construct the command
@ -76,20 +102,51 @@ struct RunEnvironment {
///
/// Will use the current virtual environment (if any) unless `isolated` is true.
/// Will create virtual environments in a temporary directory (if necessary).
fn environment_for_run(isolated: bool, cache: &Cache) -> Result<RunEnvironment> {
if !isolated {
// Return the active environment if it exists
async fn environment_for_run(
requirements: &[RequirementsSource],
isolated: bool,
cache: &Cache,
printer: Printer,
) -> Result<RunEnvironment> {
let current_venv = if isolated {
None
} else {
// Find the active environment if it exists
match PythonEnvironment::from_virtualenv(cache) {
Ok(env) => {
return Ok(RunEnvironment {
python: env,
_temp_dir_drop: None,
})
}
Err(uv_interpreter::Error::VenvNotFound) => {}
Ok(env) => Some(env),
Err(uv_interpreter::Error::VenvNotFound) => None,
Err(err) => return Err(err.into()),
};
}
};
// TODO(zanieb): Support client configuration
let client_builder = BaseClientBuilder::default();
// Read all requirements from the provided sources.
// TODO(zanieb): Consider allowing overrides and constraints
// TODO(zanieb): Allow specifying extras somehow
let spec =
RequirementsSpecification::from_simple_sources(requirements, &client_builder).await?;
// Check if the current environment satisfies the requirements
if let Some(venv) = current_venv {
// Determine the set of installed packages.
let site_packages = SitePackages::from_executable(&venv)?;
// If the requirements are already satisfied, we're done. Ideally, the resolver would be fast
// enough to let us remove this check. But right now, for large environments, it's an order of
// magnitude faster to validate the environment than to resolve the requirements.
if spec.source_trees.is_empty()
&& site_packages.satisfies(&spec.requirements, &spec.editables, &spec.constraints)?
{
debug!("Current environment satisfies requirements");
return Ok(RunEnvironment {
python: venv,
_temp_dir_drop: None,
});
}
}
// Otherwise, we need a new environment
// Find an interpreter to use
// TODO(zanieb): Populate `python` from the user
@ -100,22 +157,472 @@ fn environment_for_run(isolated: bool, cache: &Cache) -> Result<RunEnvironment>
PythonEnvironment::from_default_python(cache)?
};
// Create a virtual environment directory
// Create a virtual environment
// TODO(zanieb): Move this path derivation elsewhere
let uv_state_path = std::env::current_dir()?.join(".uv");
fs_err::create_dir_all(&uv_state_path)?;
let tmpdir = tempdir_in(uv_state_path)?;
let venv = uv_virtualenv::create_venv(
tmpdir.path(),
python_env.into_interpreter(),
uv_virtualenv::Prompt::None,
false,
Vec::new(),
)?;
// Determine the tags, markers, and interpreter to use for resolution.
let interpreter = venv.interpreter().clone();
let tags = venv.interpreter().tags()?;
let markers = venv.interpreter().markers();
// Collect the set of required hashes.
// TODO(zanieb): Support hash checking
let hasher = HashStrategy::None;
// TODO(zanieb): Support index url configs
let index_locations = IndexLocations::default();
// TODO(zanieb): Support client options e.g. offline, tls, etc.
// Initialize the registry client.
let client = RegistryClientBuilder::new(cache.clone())
.markers(markers)
.platform(interpreter.platform())
.build();
// TODO(zanieb): Consider support for find links
let flat_index = FlatIndex::default();
// TODO(zanieb): Consider support for shared builds
// Determine whether to enable build isolation.
let build_isolation = BuildIsolation::Isolated;
// TODO(zanieb): Consider no-binary and no-build support
let no_build = NoBuild::None;
let no_binary = NoBinary::None;
// Create a shared in-memory index.
let index = InMemoryIndex::default();
// Track in-flight downloads, builds, etc., across resolutions.
let in_flight = InFlight::default();
let link_mode = LinkMode::default();
let config_settings = ConfigSettings::default();
// Create a build dispatch.
let build_dispatch = BuildDispatch::new(
&client,
cache,
&interpreter,
&index_locations,
&flat_index,
&index,
&in_flight,
SetupPyStrategy::default(),
&config_settings,
build_isolation,
link_mode,
&no_build,
&no_binary,
);
// TODO(zanieb): Consider `exclude-newer` support
// Resolve the requirements from the provided sources.
let requirements = {
// Convert from unnamed to named requirements.
let mut requirements = NamedRequirementsResolver::new(
spec.requirements,
&hasher,
&build_dispatch,
&client,
&index,
)
.with_reporter(ResolverReporter::from(printer))
.resolve()
.await?;
// Resolve any source trees into requirements.
if !spec.source_trees.is_empty() {
requirements.extend(
SourceTreeResolver::new(
spec.source_trees,
&ExtrasSpecification::None,
&hasher,
&build_dispatch,
&client,
&index,
)
.with_reporter(ResolverReporter::from(printer))
.resolve()
.await?,
);
}
requirements
};
let options = OptionsBuilder::new()
// TODO(zanieb): Support resolver options
// .resolution_mode(resolution_mode)
// .prerelease_mode(prerelease_mode)
// .dependency_mode(dependency_mode)
// .exclude_newer(exclude_newer)
.build();
// Resolve the requirements.
let resolution = match resolve(
requirements,
spec.project,
&hasher,
&interpreter,
tags,
markers,
&client,
&flat_index,
&index,
&build_dispatch,
options,
printer,
)
.await
{
Ok(resolution) => Resolution::from(resolution),
Err(err) => return Err(err.into()),
};
// Re-initialize the in-flight map.
let in_flight = InFlight::default();
// Sync the environment.
install(
&resolution,
SitePackages::from_executable(&venv)?,
&no_binary,
link_mode,
&index_locations,
&hasher,
tags,
&client,
&in_flight,
&build_dispatch,
cache,
&venv,
printer,
)
.await?;
// Create the environment
// TODO(zanieb): Add dependencies to the env
Ok(RunEnvironment {
python: uv_virtualenv::create_venv(
tmpdir.path(),
python_env.into_interpreter(),
uv_virtualenv::Prompt::None,
false,
Vec::new(),
)?,
python: venv,
_temp_dir_drop: Some(tmpdir),
})
}
/// Resolve a set of requirements, similar to running `pip compile`.
#[allow(clippy::too_many_arguments)]
async fn resolve(
requirements: Vec<Requirement>,
project: Option<PackageName>,
hasher: &HashStrategy,
interpreter: &Interpreter,
tags: &Tags,
markers: &MarkerEnvironment,
client: &RegistryClient,
flat_index: &FlatIndex,
index: &InMemoryIndex,
build_dispatch: &BuildDispatch<'_>,
options: Options,
printer: Printer,
) -> Result<ResolutionGraph, Error> {
let start = std::time::Instant::now();
let exclusions = Exclusions::None;
let preferences = Vec::new();
let constraints = Constraints::default();
let overrides = Overrides::default();
let editables = Vec::new();
let installed_packages = EmptyInstalledPackages;
// Determine any lookahead requirements.
let lookaheads = LookaheadResolver::new(
&requirements,
&constraints,
&overrides,
&editables,
hasher,
build_dispatch,
client,
index,
)
.with_reporter(ResolverReporter::from(printer))
.resolve(markers)
.await?;
// Create a manifest of the requirements.
let manifest = Manifest::new(
requirements,
constraints,
overrides,
preferences,
project,
editables,
exclusions,
lookaheads,
);
// Resolve the dependencies.
let resolver = Resolver::new(
manifest,
options,
markers,
interpreter,
tags,
client,
flat_index,
index,
hasher,
build_dispatch,
&installed_packages,
)?
.with_reporter(ResolverReporter::from(printer));
let resolution = resolver.resolve().await?;
let s = if resolution.len() == 1 { "" } else { "s" };
writeln!(
printer.stderr(),
"{}",
format!(
"Resolved {} in {}",
format!("{} package{}", resolution.len(), s).bold(),
elapsed(start.elapsed())
)
.dimmed()
)?;
// Notify the user of any diagnostics.
for diagnostic in resolution.diagnostics() {
writeln!(
printer.stderr(),
"{}{} {}",
"warning".yellow().bold(),
":".bold(),
diagnostic.message().bold()
)?;
}
Ok(resolution)
}
/// Install a set of requirements into the current environment.
#[allow(clippy::too_many_arguments)]
async fn install(
resolution: &Resolution,
site_packages: SitePackages<'_>,
no_binary: &NoBinary,
link_mode: LinkMode,
index_urls: &IndexLocations,
hasher: &HashStrategy,
tags: &Tags,
client: &RegistryClient,
in_flight: &InFlight,
build_dispatch: &BuildDispatch<'_>,
cache: &Cache,
venv: &PythonEnvironment,
printer: Printer,
) -> Result<(), Error> {
let start = std::time::Instant::now();
let requirements = resolution.requirements();
// Partition into those that should be linked from the cache (`local`), those that need to be
// downloaded (`remote`), and those that should be removed (`extraneous`).
let plan = Planner::with_requirements(&requirements)
.build(
site_packages,
&Reinstall::None,
no_binary,
hasher,
index_urls,
cache,
venv,
tags,
)
.context("Failed to determine installation plan")?;
let Plan {
cached,
remote,
reinstalls,
installed: _,
extraneous: _,
} = plan;
// Nothing to do.
if remote.is_empty() && cached.is_empty() {
let s = if resolution.len() == 1 { "" } else { "s" };
writeln!(
printer.stderr(),
"{}",
format!(
"Audited {} in {}",
format!("{} package{}", resolution.len(), s).bold(),
elapsed(start.elapsed())
)
.dimmed()
)?;
return Ok(());
}
// Map any registry-based requirements back to those returned by the resolver.
let remote = remote
.iter()
.map(|dist| {
resolution
.get_remote(&dist.name)
.cloned()
.expect("Resolution should contain all packages")
})
.collect::<Vec<_>>();
// Download, build, and unzip any missing distributions.
let wheels = if remote.is_empty() {
vec![]
} else {
let start = std::time::Instant::now();
let downloader = Downloader::new(cache, tags, hasher, client, build_dispatch)
.with_reporter(DownloadReporter::from(printer).with_length(remote.len() as u64));
let wheels = downloader
.download(remote.clone(), in_flight)
.await
.context("Failed to download distributions")?;
let s = if wheels.len() == 1 { "" } else { "s" };
writeln!(
printer.stderr(),
"{}",
format!(
"Downloaded {} in {}",
format!("{} package{}", wheels.len(), s).bold(),
elapsed(start.elapsed())
)
.dimmed()
)?;
wheels
};
// Install the resolved distributions.
let wheels = wheels.into_iter().chain(cached).collect::<Vec<_>>();
if !wheels.is_empty() {
let start = std::time::Instant::now();
uv_installer::Installer::new(venv)
.with_link_mode(link_mode)
.with_reporter(InstallReporter::from(printer).with_length(wheels.len() as u64))
.install(&wheels)?;
let s = if wheels.len() == 1 { "" } else { "s" };
writeln!(
printer.stderr(),
"{}",
format!(
"Installed {} in {}",
format!("{} package{}", wheels.len(), s).bold(),
elapsed(start.elapsed())
)
.dimmed()
)?;
}
for event in reinstalls
.into_iter()
.map(|distribution| ChangeEvent {
dist: LocalDist::from(distribution),
kind: ChangeEventKind::Removed,
})
.chain(wheels.into_iter().map(|distribution| ChangeEvent {
dist: LocalDist::from(distribution),
kind: ChangeEventKind::Added,
}))
.sorted_unstable_by(|a, b| {
a.dist
.name()
.cmp(b.dist.name())
.then_with(|| a.kind.cmp(&b.kind))
.then_with(|| a.dist.installed_version().cmp(&b.dist.installed_version()))
})
{
match event.kind {
ChangeEventKind::Added => {
writeln!(
printer.stderr(),
" {} {}{}",
"+".green(),
event.dist.name().as_ref().bold(),
event.dist.installed_version().to_string().dimmed()
)?;
}
ChangeEventKind::Removed => {
writeln!(
printer.stderr(),
" {} {}{}",
"-".red(),
event.dist.name().as_ref().bold(),
event.dist.installed_version().to_string().dimmed()
)?;
}
}
}
// TODO(konstin): Also check the cache whether any cached or installed dist is already known to
// have been yanked, we currently don't show this message on the second run anymore
for dist in &remote {
let Some(file) = dist.file() else {
continue;
};
match &file.yanked {
None | Some(Yanked::Bool(false)) => {}
Some(Yanked::Bool(true)) => {
writeln!(
printer.stderr(),
"{}{} {dist} is yanked.",
"warning".yellow().bold(),
":".bold(),
)?;
}
Some(Yanked::Reason(reason)) => {
writeln!(
printer.stderr(),
"{}{} {dist} is yanked (reason: \"{reason}\").",
"warning".yellow().bold(),
":".bold(),
)?;
}
}
}
Ok(())
}
#[derive(thiserror::Error, Debug)]
enum Error {
#[error(transparent)]
Resolve(#[from] uv_resolver::ResolveError),
#[error(transparent)]
Client(#[from] uv_client::Error),
#[error(transparent)]
Platform(#[from] platform_tags::PlatformError),
#[error(transparent)]
Hash(#[from] uv_types::HashStrategyError),
#[error(transparent)]
Io(#[from] std::io::Error),
#[error(transparent)]
Fmt(#[from] std::fmt::Error),
#[error(transparent)]
Anyhow(#[from] anyhow::Error),
}

View File

@ -573,7 +573,37 @@ async fn run() -> Result<ExitStatus> {
)
.await
}
Commands::Run(args) => commands::run(args.command, args.args, args.isolated, &cache).await,
Commands::Run(args) => {
let requirements = args
.with
.into_iter()
.map(RequirementsSource::from_package)
// TODO(zanieb): Consider editable package support. What benefit do these have in an ephemeral
// environment?
// .chain(
// args.with_editable
// .into_iter()
// .map(RequirementsSource::Editable),
// )
// TODO(zanieb): Consider requirements file support, this comes with additional complexity due to
// to the extensive configuration allowed in requirements files
// .chain(
// args.with_requirements
// .into_iter()
// .map(RequirementsSource::from_requirements_file),
// )
.collect::<Vec<_>>();
commands::run(
args.command,
args.args,
&requirements,
args.isolated,
&cache,
printer,
)
.await
}
#[cfg(feature = "self-update")]
Commands::Self_(SelfNamespace {
command: SelfCommand::Update,