From ce513e134bcb6517e9a43077087915204633cf4a Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Thu, 25 Jul 2024 19:17:48 -0400 Subject: [PATCH] Use sitecustomize.py to implement environment layering --- crates/uv/src/commands/project/run.rs | 52 +++++++++++++------------- crates/uv/src/commands/tool/run.rs | 11 ------ crates/uv/tests/run.rs | 54 +++++++++++++++++++++++++++ crates/uv/tests/tool_run.rs | 2 - 4 files changed, 80 insertions(+), 39 deletions(-) diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs index 34348a8f0485..76117e368eb1 100644 --- a/crates/uv/src/commands/project/run.rs +++ b/crates/uv/src/commands/project/run.rs @@ -3,7 +3,7 @@ use std::ffi::OsString; use std::fmt::Write; use std::path::{Path, PathBuf}; -use anyhow::{bail, Context, Result}; +use anyhow::{anyhow, bail, Context, Result}; use itertools::Itertools; use owo_colors::OwoColorize; use tokio::process::Command; @@ -14,7 +14,7 @@ use uv_cache::Cache; use uv_cli::ExternalCommand; use uv_client::{BaseClientBuilder, Connectivity}; use uv_configuration::{Concurrency, ExtrasSpecification, PreviewMode}; -use uv_fs::Simplified; +use uv_fs::{PythonExt, Simplified}; use uv_installer::{SatisfiesResult, SitePackages}; use uv_normalize::PackageName; use uv_python::{ @@ -412,6 +412,30 @@ pub(crate) async fn run( } }; + // If we're running in an ephemeral environment, add a `sitecustomize.py` to enable loading of + // the base environment's site packages. Setting `PYTHONPATH` is insufficient, as it doesn't + // resolve `.pth` files in the base environment. + if let Some(ephemeral_env) = ephemeral_env.as_ref() { + if let Some(base_interpreter) = base_interpreter.as_ref() { + let ephemeral_site_packages = ephemeral_env + .site_packages() + .next() + .ok_or_else(|| anyhow!("Ephemeral environment has no site packages directory"))?; + let base_site_packages = base_interpreter + .site_packages() + .next() + .ok_or_else(|| anyhow!("Base environment has no site packages directory"))?; + + fs_err::write( + ephemeral_site_packages.join("sitecustomize.py"), + format!( + "import site; site.addsitedir(\"{}\")", + base_site_packages.escape_for_python() + ), + )?; + } + } + debug!("Running `{command}`"); let mut process = Command::from(&command); @@ -437,30 +461,6 @@ pub(crate) async fn run( )?; process.env("PATH", new_path); - // Construct the `PYTHONPATH` environment variable. - let new_python_path = std::env::join_paths( - ephemeral_env - .as_ref() - .map(PythonEnvironment::site_packages) - .into_iter() - .flatten() - .chain( - base_interpreter - .as_ref() - .map(Interpreter::site_packages) - .into_iter() - .flatten(), - ) - .map(PathBuf::from) - .chain( - std::env::var_os("PYTHONPATH") - .as_ref() - .iter() - .flat_map(std::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 diff --git a/crates/uv/src/commands/tool/run.rs b/crates/uv/src/commands/tool/run.rs index f0c335150a55..0dab15151a5a 100644 --- a/crates/uv/src/commands/tool/run.rs +++ b/crates/uv/src/commands/tool/run.rs @@ -123,17 +123,6 @@ pub(crate) async fn run( )?; process.env("PATH", new_path); - // Construct the `PYTHONPATH` environment variable. - let new_python_path = std::env::join_paths( - environment.site_packages().map(PathBuf::from).chain( - std::env::var_os("PYTHONPATH") - .as_ref() - .iter() - .flat_map(std::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 diff --git a/crates/uv/tests/run.rs b/crates/uv/tests/run.rs index 71db45856d48..a5bf5e149dac 100644 --- a/crates/uv/tests/run.rs +++ b/crates/uv/tests/run.rs @@ -769,3 +769,57 @@ fn run_requirements_txt_arguments() -> Result<()> { Ok(()) } + +/// Ensure that we can import from the root project when layering `--with` requirements. +#[test] +fn run_editable() -> Result<()> { + let context = TestContext::new("3.12"); + + let pyproject_toml = context.temp_dir.child("pyproject.toml"); + pyproject_toml.write_str(indoc! { r#" + [build-system] + requires = ["hatchling"] + build-backend = "hatchling.build" + + [project] + name = "foo" + version = "1.0.0" + requires-python = ">=3.8" + dependencies = [] + "# + })?; + + let src = context.temp_dir.child("src").child("foo"); + src.create_dir_all()?; + + let init = src.child("__init__.py"); + init.touch()?; + + let main = context.temp_dir.child("main.py"); + main.write_str(indoc! { r" + import foo + print('Hello, world!') + " + })?; + + // We treat arguments before the command as uv arguments + uv_snapshot!(context.filters(), context.run().arg("--with").arg("iniconfig").arg("main.py"), @r###" + success: true + exit_code: 0 + ----- stdout ----- + Hello, world! + + ----- stderr ----- + warning: `uv run` is experimental and may change without warning + Resolved 1 package in [TIME] + Prepared 1 package in [TIME] + Installed 1 package in [TIME] + + foo==1.0.0 (from file://[TEMP_DIR]/) + Resolved 1 package in [TIME] + Prepared 1 package in [TIME] + Installed 1 package in [TIME] + + iniconfig==2.0.0 + "###); + + Ok(()) +} diff --git a/crates/uv/tests/tool_run.rs b/crates/uv/tests/tool_run.rs index 4c5a27167a51..e19e4b51e7cb 100644 --- a/crates/uv/tests/tool_run.rs +++ b/crates/uv/tests/tool_run.rs @@ -630,7 +630,6 @@ fn tool_run_requirements_txt() { let requirements_txt = context.temp_dir.child("requirements.txt"); requirements_txt.write_str("iniconfig").unwrap(); - // We treat arguments before the command as uv arguments uv_snapshot!(context.filters(), context.tool_run() .arg("--with-requirements") .arg("requirements.txt") @@ -680,7 +679,6 @@ fn tool_run_requirements_txt_arguments() { }) .unwrap(); - // We treat arguments before the command as uv arguments uv_snapshot!(context.filters(), context.tool_run() .arg("--with-requirements") .arg("requirements.txt")