diff --git a/crates/uv-fs/src/lib.rs b/crates/uv-fs/src/lib.rs index dcc0f00b2..17f52dcf5 100644 --- a/crates/uv-fs/src/lib.rs +++ b/crates/uv-fs/src/lib.rs @@ -84,6 +84,8 @@ pub async fn read_to_string_transcode(path: impl AsRef) -> std::io::Result /// junction at the same path. /// /// Note that because junctions are used, the source must be a directory. +/// +/// Changes to this function should be reflected in [`create_symlink`]. #[cfg(windows)] pub fn replace_symlink(src: impl AsRef, dst: impl AsRef) -> std::io::Result<()> { // If the source is a file, we can't create a junction @@ -138,6 +140,38 @@ pub fn replace_symlink(src: impl AsRef, dst: impl AsRef) -> std::io: } } +/// Create a symlink at `dst` pointing to `src`. +/// +/// On Windows, this uses the `junction` crate to create a junction point. +/// +/// Note that because junctions are used, the source must be a directory. +/// +/// Changes to this function should be reflected in [`replace_symlink`]. +#[cfg(windows)] +pub fn create_symlink(src: impl AsRef, dst: impl AsRef) -> std::io::Result<()> { + // If the source is a file, we can't create a junction + if src.as_ref().is_file() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "Cannot create a junction for {}: is not a directory", + src.as_ref().display() + ), + )); + } + + junction::create( + dunce::simplified(src.as_ref()), + dunce::simplified(dst.as_ref()), + ) +} + +/// Create a symlink at `dst` pointing to `src`. +#[cfg(unix)] +pub fn create_symlink(src: impl AsRef, dst: impl AsRef) -> std::io::Result<()> { + fs_err::os::unix::fs::symlink(src.as_ref(), dst.as_ref()) +} + #[cfg(unix)] pub fn remove_symlink(path: impl AsRef) -> std::io::Result<()> { fs_err::remove_file(path.as_ref()) diff --git a/crates/uv-trampoline-builder/src/lib.rs b/crates/uv-trampoline-builder/src/lib.rs index 2e1cde872..1a25b9454 100644 --- a/crates/uv-trampoline-builder/src/lib.rs +++ b/crates/uv-trampoline-builder/src/lib.rs @@ -41,6 +41,7 @@ const MAGIC_NUMBER_SIZE: usize = 4; pub struct Launcher { pub kind: LauncherKind, pub python_path: PathBuf, + payload: Vec, } impl Launcher { @@ -109,11 +110,69 @@ impl Launcher { String::from_utf8(buffer).map_err(|err| Error::InvalidPath(err.utf8_error()))?, ); + #[allow(clippy::cast_possible_truncation)] + let file_size = { + let raw_length = file + .seek(io::SeekFrom::End(0)) + .map_err(|e| Error::InvalidLauncherSeek("size probe".into(), 0, e))?; + + if raw_length > usize::MAX as u64 { + return Err(Error::InvalidDataLength(raw_length)); + } + + // SAFETY: Above we guarantee the length is less than uszie + raw_length as usize + }; + + // Read the payload + file.seek(io::SeekFrom::Start(0)) + .map_err(|e| Error::InvalidLauncherSeek("rewind".into(), 0, e))?; + let payload_len = + file_size.saturating_sub(MAGIC_NUMBER_SIZE + PATH_LENGTH_SIZE + path_length); + let mut buffer = vec![0u8; payload_len]; + file.read_exact(&mut buffer) + .map_err(|err| Error::InvalidLauncherRead("payload".into(), err))?; + Ok(Some(Self { kind, + payload: buffer, python_path: path, })) } + + pub fn write_to_file(self, file: &mut File) -> Result<(), Error> { + let python_path = self.python_path.simplified_display().to_string(); + + if python_path.len() > MAX_PATH_LENGTH as usize { + return Err(Error::InvalidPathLength( + u32::try_from(python_path.len()).expect("path length already checked"), + )); + } + + let mut launcher: Vec = Vec::with_capacity( + self.payload.len() + python_path.len() + PATH_LENGTH_SIZE + MAGIC_NUMBER_SIZE, + ); + launcher.extend_from_slice(&self.payload); + launcher.extend_from_slice(python_path.as_bytes()); + launcher.extend_from_slice( + &u32::try_from(python_path.len()) + .expect("file path should be smaller than 4GB") + .to_le_bytes(), + ); + launcher.extend_from_slice(self.kind.magic_number()); + + file.write_all(&launcher)?; + Ok(()) + } + + #[must_use] + pub fn with_python_path(self, path: PathBuf) -> Self { + Self { + kind: self.kind, + payload: self.payload, + python_path: path, + } + } } /// The kind of trampoline launcher to create. @@ -177,6 +236,8 @@ pub enum Error { Io(#[from] io::Error), #[error("Only paths with a length up to 32KB are supported but found a length of {0} bytes")] InvalidPathLength(u32), + #[error("Only data with a length up to usize is supported but found a length of {0} bytes")] + InvalidDataLength(u64), #[error("Failed to parse executable path")] InvalidPath(#[source] Utf8Error), #[error("Failed to seek to {0} at offset {1}")] diff --git a/crates/uv/src/commands/project/environment.rs b/crates/uv/src/commands/project/environment.rs index b5bb3fd23..af3b3b351 100644 --- a/crates/uv/src/commands/project/environment.rs +++ b/crates/uv/src/commands/project/environment.rs @@ -78,6 +78,20 @@ impl EphemeralEnvironment { )?; Ok(()) } + + /// Returns the path to the environment's scripts directory. + pub(crate) fn scripts(&self) -> &Path { + self.0.scripts() + } + + /// Returns the path to the environment's Python executable. + pub(crate) fn sys_executable(&self) -> &Path { + self.0.interpreter().sys_executable() + } + + pub(crate) fn sys_prefix(&self) -> &Path { + self.0.interpreter().sys_prefix() + } } /// A [`PythonEnvironment`] stored in the cache. diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs index ba8935013..44d0dc474 100644 --- a/crates/uv/src/commands/project/run.rs +++ b/crates/uv/src/commands/project/run.rs @@ -9,8 +9,9 @@ use anyhow::{Context, anyhow, bail}; use futures::StreamExt; use itertools::Itertools; use owo_colors::OwoColorize; +use thiserror::Error; use tokio::process::Command; -use tracing::{debug, warn}; +use tracing::{debug, trace, warn}; use url::Url; use uv_cache::Cache; @@ -22,7 +23,7 @@ use uv_configuration::{ }; use uv_distribution_types::Requirement; use uv_fs::which::is_executable; -use uv_fs::{PythonExt, Simplified}; +use uv_fs::{PythonExt, Simplified, create_symlink}; use uv_installer::{SatisfiesResult, SitePackages}; use uv_normalize::{DefaultExtras, DefaultGroups, PackageName}; use uv_python::{ @@ -1071,6 +1072,67 @@ hint: If you are running a script with `{}` in the shebang, you may need to incl requirements_site_packages.escape_for_python(), ))?; + // N.B. The order here matters — earlier interpreters take precedence over the + // later ones. + for interpreter in [requirements_env.interpreter(), &base_interpreter] { + // Copy each entrypoint from the base environments to the ephemeral environment, + // updating the Python executable target to ensure they run in the ephemeral + // environment. + for entry in fs_err::read_dir(interpreter.scripts())? { + let entry = entry?; + if !entry.file_type()?.is_file() { + continue; + } + match copy_entrypoint( + &entry.path(), + &ephemeral_env.scripts().join(entry.file_name()), + interpreter.sys_executable(), + ephemeral_env.sys_executable(), + ) { + Ok(()) => {} + // If the entrypoint already exists, skip it. + Err(CopyEntrypointError::Io(err)) + if err.kind() == std::io::ErrorKind::AlreadyExists => + { + trace!( + "Skipping copy of entrypoint `{}`: already exists", + &entry.path().display() + ); + } + Err(err) => return Err(err.into()), + } + } + + // Link data directories from the base environment to the ephemeral environment. + // + // This is critical for Jupyter Lab, which cannot operate without the files it + // writes to `/share/jupyter`. + // + // See https://github.com/jupyterlab/jupyterlab/issues/17716 + for dir in &["etc/jupyter", "share/jupyter"] { + let source = interpreter.sys_prefix().join(dir); + if !matches!(source.try_exists(), Ok(true)) { + continue; + } + if !source.is_dir() { + continue; + } + let target = ephemeral_env.sys_prefix().join(dir); + if let Some(parent) = target.parent() { + fs_err::create_dir_all(parent)?; + } + match create_symlink(&source, &target) { + Ok(()) => trace!( + "Created link for {} -> {}", + target.user_display(), + source.user_display() + ), + Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {} + Err(err) => return Err(err.into()), + } + } + } + // Write the `sys.prefix` of the parent environment to the `extends-environment` key of the `pyvenv.cfg` // file. This helps out static-analysis tools such as ty (see docs on // `CachedEnvironment::set_parent_environment`). @@ -1669,3 +1731,92 @@ fn read_recursion_depth_from_environment_variable() -> anyhow::Result { .parse::() .with_context(|| format!("invalid value for {}", EnvVars::UV_RUN_RECURSION_DEPTH)) } + +#[derive(Error, Debug)] +enum CopyEntrypointError { + #[error(transparent)] + Io(#[from] std::io::Error), + #[cfg(windows)] + #[error(transparent)] + Trampoline(#[from] uv_trampoline_builder::Error), +} + +/// Create a copy of the entrypoint at `source` at `target`, if it has a Python shebang, replacing +/// the previous Python executable with a new one. +/// +/// This is a no-op if the target already exists. +/// +/// Note on Windows, the entrypoints do not use shebangs and require a rewrite of the trampoline. +#[cfg(unix)] +fn copy_entrypoint( + source: &Path, + target: &Path, + previous_executable: &Path, + python_executable: &Path, +) -> Result<(), CopyEntrypointError> { + use std::io::Write; + use std::os::unix::fs::PermissionsExt; + + use fs_err::os::unix::fs::OpenOptionsExt; + + let contents = fs_err::read_to_string(source)?; + + let Some(contents) = contents + // Check for a relative path or relocatable shebang + .strip_prefix( + r#"#!/bin/sh +'''exec' "$(dirname -- "$(realpath -- "$0")")"/'python' "$0" "$@" +' ''' +"#, + ) + // Or an absolute path shebang + .or_else(|| contents.strip_prefix(&format!("#!{}\n", previous_executable.display()))) + else { + // If it's not a Python shebang, we'll skip it + trace!( + "Skipping copy of entrypoint `{}`: does not start with expected shebang", + source.user_display() + ); + return Ok(()); + }; + + let contents = format!("#!{}\n{}", python_executable.display(), contents); + let mode = fs_err::metadata(source)?.permissions().mode(); + let mut file = fs_err::OpenOptions::new() + .create_new(true) + .write(true) + .mode(mode) + .open(target)?; + file.write_all(contents.as_bytes())?; + + trace!("Updated entrypoint at {}", target.user_display()); + + Ok(()) +} + +/// Create a copy of the entrypoint at `source` at `target`, if it's a Python script launcher, +/// replacing the target Python executable with a new one. +#[cfg(windows)] +fn copy_entrypoint( + source: &Path, + target: &Path, + _previous_executable: &Path, + python_executable: &Path, +) -> Result<(), CopyEntrypointError> { + use uv_trampoline_builder::Launcher; + + let Some(launcher) = Launcher::try_from_path(source)? else { + return Ok(()); + }; + + let launcher = launcher.with_python_path(python_executable.to_path_buf()); + let mut file = fs_err::OpenOptions::new() + .create_new(true) + .write(true) + .open(target)?; + launcher.write_to_file(&mut file)?; + + trace!("Updated entrypoint at {}", target.user_display()); + + Ok(()) +} diff --git a/crates/uv/tests/it/run.rs b/crates/uv/tests/it/run.rs index ad8672788..2e9762a60 100644 --- a/crates/uv/tests/it/run.rs +++ b/crates/uv/tests/it/run.rs @@ -1319,6 +1319,181 @@ fn run_with_pyvenv_cfg_file() -> Result<()> { Ok(()) } +#[test] +fn run_with_overlay_interpreter() -> Result<()> { + let context = TestContext::new("3.12").with_filtered_exe_suffix(); + + let pyproject_toml = context.temp_dir.child("pyproject.toml"); + pyproject_toml.write_str(indoc! { r#" + [project] + name = "foo" + version = "1.0.0" + requires-python = ">=3.8" + dependencies = ["anyio"] + + [build-system] + requires = ["setuptools>=42"] + build-backend = "setuptools.build_meta" + + [project.scripts] + main = "foo:main" + "# + })?; + + let foo = context.temp_dir.child("src").child("foo"); + foo.create_dir_all()?; + let init_py = foo.child("__init__.py"); + init_py.write_str(indoc! { r#" + import sys + import shutil + from pathlib import Path + + def show_python(): + print(sys.executable) + + def copy_entrypoint(): + base = Path(sys.executable) + shutil.copyfile(base.with_name("main").with_suffix(base.suffix), sys.argv[1]) + + def main(): + show_python() + if len(sys.argv) > 1: + copy_entrypoint() + "# + })?; + + // The project's entrypoint should be rewritten to use the overlay interpreter. + uv_snapshot!(context.filters(), context.run().arg("--with").arg("iniconfig").arg("main").arg(context.temp_dir.child("main").as_os_str()), @r" + success: true + exit_code: 0 + ----- stdout ----- + [CACHE_DIR]/builds-v0/[TMP]/python + + ----- stderr ----- + Resolved 6 packages in [TIME] + Prepared 4 packages in [TIME] + Installed 4 packages in [TIME] + + anyio==4.3.0 + + foo==1.0.0 (from file://[TEMP_DIR]/) + + idna==3.6 + + sniffio==1.3.1 + Resolved 1 package in [TIME] + Prepared 1 package in [TIME] + Installed 1 package in [TIME] + + iniconfig==2.0.0 + "); + + #[cfg(unix)] + insta::with_settings!({ + filters => context.filters(), + }, { + assert_snapshot!( + context.read("main"), @r##" + #![CACHE_DIR]/builds-v0/[TMP]/python + # -*- coding: utf-8 -*- + import sys + from foo import main + if __name__ == "__main__": + if sys.argv[0].endswith("-script.pyw"): + sys.argv[0] = sys.argv[0][:-11] + elif sys.argv[0].endswith(".exe"): + sys.argv[0] = sys.argv[0][:-4] + sys.exit(main()) + "## + ); + } + ); + + // The package, its dependencies, and the overlay dependencies should be available. + context + .run() + .arg("--with") + .arg("iniconfig") + .arg("python") + .arg("-c") + .arg("import foo; import anyio; import iniconfig") + .assert() + .success(); + + // When layering the project on top (via `--with`), the overlay interpreter also should be used. + uv_snapshot!(context.filters(), context.run().arg("--no-project").arg("--with").arg(".").arg("main"), @r" + success: true + exit_code: 0 + ----- stdout ----- + [CACHE_DIR]/builds-v0/[TMP]/python + + ----- stderr ----- + Resolved 4 packages in [TIME] + Prepared 1 package in [TIME] + Installed 4 packages in [TIME] + + anyio==4.3.0 + + foo==1.0.0 (from file://[TEMP_DIR]/) + + idna==3.6 + + sniffio==1.3.1 + "); + + // Switch to a relocatable virtual environment. + context.venv().arg("--relocatable").assert().success(); + + // The project's entrypoint should be rewritten to use the overlay interpreter. + uv_snapshot!(context.filters(), context.run().arg("--with").arg("iniconfig").arg("main").arg(context.temp_dir.child("main").as_os_str()), @r" + success: true + exit_code: 0 + ----- stdout ----- + [CACHE_DIR]/builds-v0/[TMP]/python + + ----- stderr ----- + Resolved 6 packages in [TIME] + Audited 4 packages in [TIME] + Resolved 1 package in [TIME] + "); + + // The package, its dependencies, and the overlay dependencies should be available. + context + .run() + .arg("--with") + .arg("iniconfig") + .arg("python") + .arg("-c") + .arg("import foo; import anyio; import iniconfig") + .assert() + .success(); + + #[cfg(unix)] + insta::with_settings!({ + filters => context.filters(), + }, { + assert_snapshot!( + context.read("main"), @r##" + #![CACHE_DIR]/builds-v0/[TMP]/python + # -*- coding: utf-8 -*- + import sys + from foo import main + if __name__ == "__main__": + if sys.argv[0].endswith("-script.pyw"): + sys.argv[0] = sys.argv[0][:-11] + elif sys.argv[0].endswith(".exe"): + sys.argv[0] = sys.argv[0][:-4] + sys.exit(main()) + "## + ); + } + ); + + // When layering the project on top (via `--with`), the overlay interpreter also should be used. + uv_snapshot!(context.filters(), context.run().arg("--no-project").arg("--with").arg(".").arg("main"), @r" + success: true + exit_code: 0 + ----- stdout ----- + [CACHE_DIR]/builds-v0/[TMP]/python + + ----- stderr ----- + Resolved 4 packages in [TIME] + "); + + Ok(()) +} + #[test] fn run_with_build_constraints() -> Result<()> { let context = TestContext::new("3.9");