Use environment layering for `uv run --with` (#3447)

## Summary

This PR takes a different approach to `--with` for `uv run`. Now,
instead of merging the requirements and re-resolving, we have two
phases: (1) sync the workspace requirements to the workspace
environment; then (2) sync the ephemeral `--with` requirements to an
ephemeral environment. The two environments are then layered by setting
the `PATH` and `PYTHONPATH` variables appropriately.

I think this approach simplifies a few things:

1. Once we have a lockfile, the semantics are much clearer, and we can
actually reuse it for the workspace. If we had to add arbitrary
dependencies via `--with`, then it's not really clear how the lockfile
would/should behave.
2. Caching becomes simpler, because we can just cache the ephemeral
environment based on the requirements.

The current version of this PR loses a few behaviors though that I need
to restore:

- `--python` support -- but I'm not yet sure how this is supposed to
behave within projects? It's also left unclear in `uv sync` and `uv
lock`.
- The "reuse the workspace environment if it already satisfies the
ephemeral requirements" behavior.

Closes #3411.

crates/uv/src/commands/workspace/run.rs

@@ -1,11 +1,10 @@
+use std::env;
 use std::ffi::OsString;
 use std::path::PathBuf;
-use std::{env, iter};
 
 use anyhow::Result;
 use itertools::Itertools;
-use owo_colors::OwoColorize;
+use tempfile::tempdir_in;
-use tempfile::{tempdir_in, TempDir};
 use tokio::process::Command;
 use tracing::debug;
 
@@ -15,7 +14,6 @@ use uv_cache::Cache;
 use uv_client::{BaseClientBuilder, RegistryClientBuilder};
 use uv_configuration::{ConfigSettings, NoBinary, NoBuild, PreviewMode, SetupPyStrategy};
 use uv_dispatch::BuildDispatch;
-use uv_fs::Simplified;
 use uv_installer::{SatisfiesResult, SitePackages};
 use uv_interpreter::PythonEnvironment;
 use uv_requirements::{ExtrasSpecification, RequirementsSource, RequirementsSpecification};
@@ -31,7 +29,7 @@ use crate::printer::Printer;
 pub(crate) async fn run(
     target: Option<String>,
     mut args: Vec<OsString>,
-    mut requirements: Vec<RequirementsSource>,
+    requirements: Vec<RequirementsSource>,
     python: Option<String>,
     isolated: bool,
     preview: PreviewMode,
@@ -58,53 +56,112 @@ pub(crate) async fn run(
         "python".to_string()
     };
 
-    // Copy the requirements into a set of overrides; we'll use this to prioritize
+    // Discover and sync the workspace.
-    // requested requirements over those discovered in the project.
+    let workspace_env = if isolated {
-    // We must retain these requirements as direct dependencies too, as overrides
+        None
-    // cannot be applied to transitive dependencies.
+    } else {
-    let overrides = requirements.clone();
+        debug!("Syncing workspace environment.");
 
-    if !isolated {
+        let Some(workspace_requirements) = workspace::find_workspace()? else {
-        if let Some(workspace_requirements) = workspace::find_workspace()? {
+            return Err(anyhow::anyhow!(
-            requirements.extend(workspace_requirements);
+                "Unable to find `pyproject.toml` for project workspace."
-        }
+            ));
-    }
+        };
 
-    // Detect the current Python interpreter.
+        let venv = PythonEnvironment::from_virtualenv(cache)?;
-    // TODO(zanieb): Create ephemeral environments
+
-    // TODO(zanieb): Accept `--python`
+        // Install the workspace requirements.
-    let run_env = environment_for_run(
+        Some(update_environment(venv, &workspace_requirements, preview, cache, printer).await?)
-        &requirements,
+    };
-        &overrides,
+
-        python.as_deref(),
+    // If necessary, create an environment for the ephemeral requirements.
-        isolated,
+    let tmpdir;
-        preview,
+    let ephemeral_env = if requirements.is_empty() {
-        cache,
+        None
-        printer,
+    } else {
-    )
+        debug!("Syncing ephemeral environment.");
-    .await?;
+
-    let python_env = run_env.python;
+        // Discover an interpreter.
+        let interpreter = if let Some(workspace_env) = &workspace_env {
+            workspace_env.interpreter().clone()
+        } else if let Some(python) = python.as_ref() {
+            PythonEnvironment::from_requested_python(python, cache)?.into_interpreter()
+        } else {
+            PythonEnvironment::from_default_python(cache)?.into_interpreter()
+        };
+
+        // TODO(charlie): If the environment satisfies the requirements, skip creation.
+        // TODO(charlie): Pass the already-installed versions as preferences, or even as the
+        // "installed" packages, so that we can skip re-installing them in the ephemeral
+        // environment.
+
+        // Create a virtual environment
+        // TODO(zanieb): Move this path derivation elsewhere
+        let uv_state_path = env::current_dir()?.join(".uv");
+        fs_err::create_dir_all(&uv_state_path)?;
+        tmpdir = tempdir_in(uv_state_path)?;
+        let venv = uv_virtualenv::create_venv(
+            tmpdir.path(),
+            interpreter,
+            uv_virtualenv::Prompt::None,
+            false,
+            false,
+        )?;
+
+        // Install the ephemeral requirements.
+        Some(update_environment(venv, &requirements, preview, cache, printer).await?)
+    };
 
     // Construct the command
     let mut process = Command::new(&command);
     process.args(&args);
 
-    // Set up the PATH
+    // Construct the `PATH` environment variable.
-    debug!(
+    let new_path = env::join_paths(
-        "Using Python {} environment at {}",
+        ephemeral_env
-        python_env.interpreter().python_version(),
+            .as_ref()
-        python_env.python_executable().user_display().cyan()
+            .map(PythonEnvironment::scripts)
-    );
+            .into_iter()
-    let new_path = if let Some(path) = env::var_os("PATH") {
+            .chain(
-        let python_env_path =
+                workspace_env
-            iter::once(python_env.scripts().to_path_buf()).chain(env::split_paths(&path));
+                    .as_ref()
-        env::join_paths(python_env_path)?
+                    .map(PythonEnvironment::scripts)
-    } else {
+                    .into_iter(),
-        OsString::from(python_env.scripts())
+            )
-    };
+            .map(PathBuf::from)
-
+            .chain(
+                env::var_os("PATH")
+                    .as_ref()
+                    .iter()
+                    .flat_map(env::split_paths),
+            ),
+    )?;
     process.env("PATH", new_path);
 
+    // Construct the `PYTHONPATH` environment variable.
+    let new_python_path = env::join_paths(
+        ephemeral_env
+            .as_ref()
+            .map(PythonEnvironment::site_packages)
+            .into_iter()
+            .flatten()
+            .chain(
+                workspace_env
+                    .as_ref()
+                    .map(PythonEnvironment::site_packages)
+                    .into_iter()
+                    .flatten(),
+            )
+            .map(PathBuf::from)
+            .chain(
+                env::var_os("PYTHONPATH")
+                    .as_ref()
+                    .iter()
+                    .flat_map(env::split_paths),
+            ),
+    )?;
+    process.env("PYTHONPATH", new_python_path);
+
     // Spawn and wait for completion
     // Standard input, output, and error streams are all inherited
     // TODO(zanieb): Throw a nicer error message if the command is not found
@@ -125,40 +182,14 @@ pub(crate) async fn run(
     }
 }
 
-struct RunEnvironment {
+/// Update a [`PythonEnvironment`] to satisfy a set of [`RequirementsSource`]s.
-    /// The Python environment to execute the run in.
+async fn update_environment(
-    python: PythonEnvironment,
+    venv: PythonEnvironment,
-    /// A temporary directory, if a new virtual environment was created.
-    ///
-    /// Included to ensure that the temporary directory exists for the length of the operation, but
-    /// is dropped at the end as appropriate.
-    _temp_dir_drop: Option<TempDir>,
-}
-
-/// Returns an environment for a `run` invocation.
-///
-/// Will use the current virtual environment (if any) unless `isolated` is true.
-/// Will create virtual environments in a temporary directory (if necessary).
-async fn environment_for_run(
     requirements: &[RequirementsSource],
-    overrides: &[RequirementsSource],
-    python: Option<&str>,
-    isolated: bool,
     preview: PreviewMode,
     cache: &Cache,
     printer: Printer,
-) -> Result<RunEnvironment> {
+) -> Result<PythonEnvironment> {
-    let current_venv = if isolated {
-        None
-    } else {
-        // Find the active environment if it exists
-        match PythonEnvironment::from_virtualenv(cache) {
-            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();
 
@@ -168,79 +199,41 @@ async fn environment_for_run(
     let spec = RequirementsSpecification::from_sources(
         requirements,
         &[],
-        overrides,
+        &[],
         &ExtrasSpecification::None,
         &client_builder,
         preview,
     )
     .await?;
 
-    // Determine an interpreter to use
-    let python_env = if let Some(python) = python {
-        PythonEnvironment::from_requested_python(python, cache)?
-    } else {
-        PythonEnvironment::from_default_python(cache)?
-    };
-
     // Check if the current environment satisfies the requirements
-    if let Some(venv) = current_venv {
+    let site_packages = SitePackages::from_executable(&venv)?;
-        // Ensure it matches the selected interpreter
-        // TODO(zanieb): We should check if a version was requested and see if the environment meets that
-        //               too but this can wait until we refactor interpreter discovery
-        if venv.root() == python_env.root() {
-            // 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
+    // If the requirements are already satisfied, we're done.
-            // enough to let us remove this check. But right now, for large environments, it's an order of
+    if spec.source_trees.is_empty() {
-            // magnitude faster to validate the environment than to resolve the requirements.
+        match site_packages.satisfies(&spec.requirements, &spec.editables, &spec.constraints)? {
-            if spec.source_trees.is_empty() {
+            SatisfiesResult::Fresh {
-                match site_packages.satisfies(
+                recursive_requirements,
-                    &spec.requirements,
+            } => {
-                    &spec.editables,
+                debug!(
-                    &spec.constraints,
+                    "All requirements satisfied: {}",
-                )? {
+                    recursive_requirements
-                    SatisfiesResult::Fresh {
+                        .iter()
-                        recursive_requirements,
+                        .map(|entry| entry.requirement.to_string())
-                    } => {
+                        .sorted()
-                        debug!(
+                        .join(" | ")
-                            "All requirements satisfied: {}",
+                );
-                            recursive_requirements
+                debug!(
-                                .iter()
+                    "All editables satisfied: {}",
-                                .map(|entry| entry.requirement.to_string())
+                    spec.editables.iter().map(ToString::to_string).join(", ")
-                                .sorted()
+                );
-                                .join(" | ")
+                return Ok(venv);
-                        );
+            }
-                        debug!(
+            SatisfiesResult::Unsatisfied(requirement) => {
-                            "All editables satisfied: {}",
+                debug!("At least one requirement is not satisfied: {requirement}");
-                            spec.editables.iter().map(ToString::to_string).join(", ")
-                        );
-                        return Ok(RunEnvironment {
-                            python: venv,
-                            _temp_dir_drop: None,
-                        });
-                    }
-                    SatisfiesResult::Unsatisfied(requirement) => {
-                        debug!("At least one requirement is not satisfied: {requirement}");
-                    }
-                }
             }
         }
     }
-    // Otherwise, we need a new environment
-
-    // Create a virtual environment
-    // TODO(zanieb): Move this path derivation elsewhere
-    let uv_state_path = 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,
-        false,
-    )?;
 
     // Determine the tags, markers, and interpreter to use for resolution.
     let interpreter = venv.interpreter().clone();
@@ -333,8 +326,5 @@ async fn environment_for_run(
     )
     .await?;
 
-    Ok(RunEnvironment {
+    Ok(venv)
-        python: venv,
-        _temp_dir_drop: Some(tmpdir),
-    })
 }

crates/uv/tests/pip_compile.rs

@@ -6400,26 +6400,19 @@ fn no_stream() -> Result<()> {
     requirements_in
         .write_str("hashb_foxglove_protocolbuffers_python==25.3.0.1.20240226043130+465630478360")?;
     let constraints_in = context.temp_dir.child("constraints.in");
-    constraints_in.write_str("protobuf<=5.26.0")?;
+    constraints_in.write_str("protobuf==5.26.0")?;
 
-    uv_snapshot!(Command::new(get_bin())
+    uv_snapshot!(context.compile_without_exclude_newer()
-        .arg("pip")
-        .arg("compile")
         .arg("requirements.in")
         .arg("-c")
         .arg("constraints.in")
         .arg("--extra-index-url")
-        .arg("https://buf.build/gen/python")
+        .arg("https://buf.build/gen/python"), @r###"
-        .arg("--cache-dir")
-        .arg(context.cache_dir.path())
-        .env("VIRTUAL_ENV", context.venv.as_os_str())
-        .env("UV_NO_WRAP", "1")
-        .current_dir(&context.temp_dir), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
     # This file was autogenerated by uv via the following command:
-    #    uv pip compile requirements.in -c constraints.in --cache-dir [CACHE_DIR]
+    #    uv pip compile --cache-dir [CACHE_DIR] requirements.in -c constraints.in
     hashb-foxglove-protocolbuffers-python==25.3.0.1.20240226043130+465630478360
     protobuf==5.26.0
         # via hashb-foxglove-protocolbuffers-python