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.

crates/uv/src/cli.rs

@@ -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)]

crates/uv/src/commands/run.rs

@@ -1,29 +1,55 @@
-use std::ffi::OsString;
+use crate::commands::reporters::{DownloadReporter, InstallReporter, ResolverReporter};
-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::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> {
+async fn environment_for_run(
-    if !isolated {
+    requirements: &[RequirementsSource],
-        // Return the active environment if it exists
+    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) => {
+            Ok(env) => Some(env),
-                return Ok(RunEnvironment {
+            Err(uv_interpreter::Error::VenvNotFound) => None,
-                    python: env,
-                    _temp_dir_drop: None,
-                })
-            }
-            Err(uv_interpreter::Error::VenvNotFound) => {}
             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(
+        python: venv,
-            tmpdir.path(),
-            python_env.into_interpreter(),
-            uv_virtualenv::Prompt::None,
-            false,
-            Vec::new(),
-        )?,
         _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),
+}

crates/uv/src/main.rs

@@ -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,