Copy entry points and Jupyter data directories into ephemeral environments

crates/uv-fs/src/lib.rs

@@ -138,6 +138,36 @@ pub fn replace_symlink(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> 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.
+#[cfg(windows)]
+pub fn create_symlink(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> 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<Path>, dst: impl AsRef<Path>) -> std::io::Result<()> {
+    fs_err::os::unix::fs::symlink(src.as_ref(), dst.as_ref())
+}
+
 #[cfg(unix)]
 pub fn remove_symlink(path: impl AsRef<Path>) -> std::io::Result<()> {
     fs_err::remove_file(path.as_ref())

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<u8>,
 }
 
 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 to_file(self, path: &Path) -> 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<u8> = 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());
+
+        fs_err::write(path, 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}")]

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.

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

@@ -10,7 +10,7 @@ use futures::StreamExt;
 use itertools::Itertools;
 use owo_colors::OwoColorize;
 use tokio::process::Command;
-use tracing::{debug, warn};
+use tracing::{debug, trace, warn};
 use url::Url;
 
 use uv_cache::Cache;
@@ -22,7 +22,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, symlink_or_copy_file};
 use uv_installer::{SatisfiesResult, SitePackages};
 use uv_normalize::{DefaultExtras, DefaultGroups, PackageName};
 use uv_python::{
@@ -1071,6 +1071,75 @@ 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;
+                    }
+
+                    copy_entrypoint(
+                        &entry.path(),
+                        &ephemeral_env.scripts().join(entry.file_name()),
+                        interpreter.sys_executable(),
+                        ephemeral_env.sys_executable(),
+                    )?;
+                }
+
+                // 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 `<prefix>/share/jupyter`.
+                //
+                // See https://github.com/jupyterlab/jupyterlab/issues/17716
+                //
+                // We only perform a shallow merge here, so if a directory occurs more than once,
+                // any children of that directory will not be merged across base interpreters.
+                for dir in &["etc/jupyter", "share/jupyter"] {
+                    let entries = match fs_err::read_dir(interpreter.sys_prefix().join(dir)) {
+                        Ok(entries) => entries,
+                        // Skip missing directories
+                        Err(err) if err.kind() == std::io::ErrorKind::NotFound => continue,
+                        Err(err) => return Err(err.into()),
+                    };
+                    fs_err::create_dir_all(ephemeral_env.sys_prefix().join(dir))?;
+                    for entry in entries {
+                        let entry = entry?;
+                        let target = ephemeral_env.sys_prefix().join(dir).join(entry.file_name());
+
+                        if entry.file_type()?.is_file() {
+                            match symlink_or_copy_file(entry.path(), &target) {
+                                Ok(()) => trace!(
+                                    "Created link for {} -> {}",
+                                    target.user_display(),
+                                    entry.path().user_display()
+                                ),
+                                Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {}
+                                Err(err) => return Err(err.into()),
+                            }
+                        } else if entry.file_type()?.is_dir() {
+                            match create_symlink(entry.path(), &target) {
+                                Ok(()) => trace!(
+                                    "Created link for {} -> {}",
+                                    target.user_display(),
+                                    entry.path().user_display()
+                                ),
+                                Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {}
+                                Err(err) => return Err(err.into()),
+                            }
+                        } else {
+                            trace!("Skipping link of entry: {}", entry.path().user_display());
+                        }
+                    }
+                }
+            }
+
             // 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`).
@@ -1095,6 +1164,10 @@ hint: If you are running a script with `{}` in the shebang, you may need to incl
     // Cast to `PythonEnvironment`.
     let ephemeral_env = ephemeral_env.map(PythonEnvironment::from);
 
+    if let Some(e) = ephemeral_env.as_ref() {
+        debug!("Using ephemeral environment at: {}", e.scripts().display());
+    }
+
     // Determine the Python interpreter to use for the command, if necessary.
     let interpreter = ephemeral_env
         .as_ref()
@@ -1669,3 +1742,80 @@ fn read_recursion_depth_from_environment_variable() -> anyhow::Result<u32> {
         .parse::<u32>()
         .with_context(|| format!("invalid value for {}", EnvVars::UV_RUN_RECURSION_DEPTH))
 }
+
+/// 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.
+///
+/// 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,
+) -> anyhow::Result<()> {
+    use std::os::unix::fs::PermissionsExt;
+
+    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!("#!{}", previous_executable.display())))
+    else {
+        // If it's not a Python shebang, we'll skip it
+        trace!(
+            "Skipping copy of entrypoint at {}: does not start with expected shebang",
+            source.user_display()
+        );
+        return Ok(());
+    };
+
+    let contents = format!("#!{}\n{}", python_executable.display(), contents);
+    fs_err::write(target, &contents)?;
+
+    match fs_err::write(target, &contents) {
+        Ok(()) => trace!("Updated entrypoint at {}", target.user_display()),
+        Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {}
+        Err(err) => return Err(err.into()),
+    }
+
+    let mut perms = fs_err::metadata(target)?.permissions();
+    perms.set_mode(0o755);
+    fs_err::set_permissions(target, perms)?;
+
+    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,
+) -> anyhow::Result<()> {
+    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());
+
+    match launcher.to_file(target) {
+        Ok(()) => trace!("Updated entrypoint at {}", target.user_display()),
+        Err(uv_trampoline_builder::Error::Io(err))
+            if err.kind() == std::io::ErrorKind::AlreadyExists => {}
+        Err(err) => return Err(err.into()),
+    }
+
+    Ok(())
+}