Add relocatable installs to support concurrency-safe cached environments

crates/install-wheel-rs/src/lib.rs

@@ -1,7 +1,6 @@
 //! Takes a wheel and installs it into a venv.
 
 use std::io;
-
 use std::path::PathBuf;
 
 use platform_info::PlatformInfoError;
@@ -34,6 +33,11 @@ pub struct Layout {
     pub os_name: String,
     /// The [`Scheme`] paths for the interpreter.
     pub scheme: Scheme,
+    /// Whether the environment is "relocatable". When enabled, any paths to the environment (e.g.,
+    /// those encoded in entrypoints and scripts) will be expressed in relative terms. As a result,
+    /// the entrypoints and scripts themselves will _not_ be relocatable, but the environment as a
+    /// whole will be.
+    pub relocatable: bool,
 }
 
 /// Note: The caller is responsible for adding the path of the wheel we're installing.

crates/install-wheel-rs/src/wheel.rs

@@ -1,14 +1,14 @@
-use std::collections::HashMap;
-use std::io::{BufRead, BufReader, Cursor, Read, Write};
-use std::path::{Path, PathBuf};
-use std::{env, io};
-
 use data_encoding::BASE64URL_NOPAD;
 use fs_err as fs;
 use fs_err::{DirEntry, File};
 use mailparse::MailHeaderMap;
 use rustc_hash::FxHashMap;
 use sha2::{Digest, Sha256};
+use std::borrow::Cow;
+use std::collections::HashMap;
+use std::io::{BufRead, BufReader, Cursor, Read, Write};
+use std::path::{Path, PathBuf};
+use std::{env, io};
 use tracing::{instrument, warn};
 use walkdir::WalkDir;
 use zip::write::FileOptions;
@@ -122,15 +122,22 @@ fn copy_and_hash(reader: &mut impl Read, writer: &mut impl Write) -> io::Result<
     ))
 }
 
-/// Format the shebang for a given Python executable.
+/// Format the shebang for a given Python executable, assuming an absolute path.
 ///
 /// Like pip, if a shebang is non-simple (too long or contains spaces), we use `/bin/sh` as the
 /// executable.
 ///
 /// See: <https://github.com/pypa/pip/blob/0ad4c94be74cc24874c6feb5bb3c2152c398a18e/src/pip/_vendor/distlib/scripts.py#L136-L165>
-fn format_shebang(executable: impl AsRef<Path>, os_name: &str) -> String {
+fn format_absolute_shebang(executable: impl AsRef<Path>, os_name: &str) -> String {
+    let executable = executable.as_ref();
+    debug_assert!(
+        executable.is_absolute(),
+        "Path must be absolute: {}",
+        executable.display()
+    );
+
     // Convert the executable to a simplified path.
-    let executable = executable.as_ref().simplified_display().to_string();
+    let executable = executable.simplified_display().to_string();
 
     // Validate the shebang.
     if os_name == "posix" {
@@ -151,6 +158,32 @@ fn format_shebang(executable: impl AsRef<Path>, os_name: &str) -> String {
     format!("#!{executable}")
 }
 
+/// Format the shebang for a given Python executable, assuming a relative path.
+///
+/// Like pip, if a shebang is non-simple (too long or contains spaces), we use `/bin/sh` as the
+/// executable.
+///
+/// See: <https://github.com/pypa/pip/blob/0ad4c94be74cc24874c6feb5bb3c2152c398a18e/src/pip/_vendor/distlib/scripts.py#L136-L165>
+fn format_relative_shebang(executable: impl AsRef<Path>, os_name: &str) -> String {
+    let executable = executable.as_ref();
+    debug_assert!(
+        executable.is_relative(),
+        "Path must be relative: {}",
+        executable.display()
+    );
+
+    // Convert the executable to a simplified path.
+    let executable = executable.simplified_display().to_string();
+
+    // Wrap in `dirname`. We assume that the relative path is fairly simple, since we know it's a
+    // relative path within a virtual environment, and so we shouldn't need to handle quotes,e tc.
+    if os_name == "posix" {
+        return format!("#!/bin/sh\n'''exec' \"$(dirname $0)/{executable}\" \"$0\" \"$@\"\n' '''");
+    }
+
+    format!("#!{executable}")
+}
+
 /// A Windows script is a minimal .exe launcher binary with the python entrypoint script appended as
 /// stored zip file.
 ///
@@ -291,12 +324,32 @@ pub(crate) fn write_script_entrypoints(
                 ))
             })?;
 
+        let python_executable = if layout.relocatable {
+            Cow::Owned(
+                pathdiff::diff_paths(&layout.sys_executable, &layout.scheme.scripts).ok_or_else(
+                    || {
+                        Error::Io(io::Error::new(
+                            io::ErrorKind::Other,
+                            format!(
+                                "Could not find relative path for: {}",
+                                layout.sys_executable.simplified_display()
+                            ),
+                        ))
+                    },
+                )?,
+            )
+        } else {
+            Cow::Borrowed(&layout.sys_executable)
+        };
+
         // Generate the launcher script.
-        let launcher_executable = get_script_executable(&layout.sys_executable, is_gui);
-        let launcher_python_script = get_script_launcher(
-            entrypoint,
-            &format_shebang(&launcher_executable, &layout.os_name),
-        );
+        let launcher_executable = get_script_executable(&python_executable, is_gui);
+        let shebang = if layout.relocatable {
+            format_relative_shebang(&launcher_executable, &layout.os_name)
+        } else {
+            format_absolute_shebang(&launcher_executable, &layout.os_name)
+        };
+        let launcher_python_script = get_script_launcher(entrypoint, &shebang);
 
         // If necessary, wrap the launcher script in a Windows launcher binary.
         if cfg!(windows) {
@@ -432,9 +485,9 @@ pub(crate) fn move_folder_recorded(
     Ok(())
 }
 
-/// Installs a single script (not an entrypoint)
+/// Installs a single script (not an entrypoint).
 ///
-/// Has to deal with both binaries files (just move) and scripts (rewrite the shebang if applicable)
+/// Has to deal with both binaries files (just move) and scripts (rewrite the shebang if applicable).
 fn install_script(
     layout: &Layout,
     site_packages: &Path,
@@ -494,7 +547,25 @@ fn install_script(
     let mut start = vec![0; placeholder_python.len()];
     script.read_exact(&mut start)?;
     let size_and_encoded_hash = if start == placeholder_python {
-        let start = format_shebang(&layout.sys_executable, &layout.os_name)
+        let python_executable = if layout.relocatable {
+            Cow::Owned(
+                pathdiff::diff_paths(&layout.sys_executable, &layout.scheme.scripts).ok_or_else(
+                    || {
+                        Error::Io(io::Error::new(
+                            io::ErrorKind::Other,
+                            format!(
+                                "Could not find relative path for: {}",
+                                layout.sys_executable.simplified_display()
+                            ),
+                        ))
+                    },
+                )?,
+            )
+        } else {
+            Cow::Borrowed(&layout.sys_executable)
+        };
+
+        let start = format_absolute_shebang(python_executable.as_ref(), &layout.os_name)
             .as_bytes()
             .to_vec();
 
@@ -779,7 +850,7 @@ mod test {
     use assert_fs::prelude::*;
     use indoc::{formatdoc, indoc};
 
-    use crate::wheel::format_shebang;
+    use crate::wheel::format_absolute_shebang;
     use crate::Error;
 
     use super::{
@@ -884,37 +955,44 @@ mod test {
     }
 
     #[test]
-    fn test_shebang() {
+    #[cfg(not(windows))]
+    fn test_absolute() {
         // By default, use a simple shebang.
         let executable = Path::new("/usr/bin/python3");
         let os_name = "posix";
-        assert_eq!(format_shebang(executable, os_name), "#!/usr/bin/python3");
+        assert_eq!(
+            format_absolute_shebang(executable, os_name),
+            "#!/usr/bin/python3"
+        );
 
         // If the path contains spaces, we should use the `exec` trick.
         let executable = Path::new("/usr/bin/path to python3");
         let os_name = "posix";
         assert_eq!(
-            format_shebang(executable, os_name),
+            format_absolute_shebang(executable, os_name),
             "#!/bin/sh\n'''exec' '/usr/bin/path to python3' \"$0\" \"$@\"\n' '''"
         );
 
         // Except on Windows...
         let executable = Path::new("/usr/bin/path to python3");
         let os_name = "nt";
         assert_eq!(
-            format_shebang(executable, os_name),
+            format_absolute_shebang(executable, os_name),
             "#!/usr/bin/path to python3"
         );
 
         // Quotes, however, are ok.
         let executable = Path::new("/usr/bin/'python3'");
         let os_name = "posix";
-        assert_eq!(format_shebang(executable, os_name), "#!/usr/bin/'python3'");
+        assert_eq!(
+            format_absolute_shebang(executable, os_name),
+            "#!/usr/bin/'python3'"
+        );
 
         // If the path is too long, we should not use the `exec` trick.
         let executable = Path::new("/usr/bin/path/to/a/very/long/executable/executable/executable/executable/executable/executable/executable/executable/name/python3");
         let os_name = "posix";
-        assert_eq!(format_shebang(executable, os_name), "#!/bin/sh\n'''exec' '/usr/bin/path/to/a/very/long/executable/executable/executable/executable/executable/executable/executable/executable/name/python3' \"$0\" \"$@\"\n' '''");
+        assert_eq!(format_absolute_shebang(executable, os_name), "#!/bin/sh\n'''exec' '/usr/bin/path/to/a/very/long/executable/executable/executable/executable/executable/executable/executable/executable/name/python3' \"$0\" \"$@\"\n' '''");
     }
 
     #[test]

crates/uv-installer/src/installer.rs

@@ -1,8 +1,7 @@
-use std::convert;
-
 use anyhow::{Context, Error, Result};
 use install_wheel_rs::{linker::LinkMode, Layout};
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
+use std::convert;
 use tokio::sync::oneshot;
 use tracing::instrument;
 

crates/uv-python/src/environment.rs

@@ -144,6 +144,16 @@ impl PythonEnvironment {
         })))
     }
 
+    /// Create a [`PythonEnvironment`] from an existing [`Interpreter`] with relocatable paths.
+    #[must_use]
+    pub fn with_relocatable(self) -> Self {
+        let inner = Arc::unwrap_or_clone(self.0);
+        Self(Arc::new(PythonEnvironmentShared {
+            interpreter: inner.interpreter.with_relocatable(),
+            ..inner
+        }))
+    }
+
     /// Returns the root (i.e., `prefix`) of the Python interpreter.
     pub fn root(&self) -> &Path {
         &self.0.root

crates/uv-python/src/interpreter.rs

@@ -45,6 +45,7 @@ pub struct Interpreter {
     prefix: Option<Prefix>,
     pointer_size: PointerSize,
     gil_disabled: bool,
+    relocatable: bool,
 }
 
 impl Interpreter {
@@ -75,6 +76,7 @@ impl Interpreter {
             tags: OnceLock::new(),
             target: None,
             prefix: None,
+            relocatable: false,
         })
     }
 
@@ -109,6 +111,7 @@ impl Interpreter {
             prefix: None,
             pointer_size: PointerSize::_64,
             gil_disabled: false,
+            relocatable: false,
         }
     }
 
@@ -143,6 +146,15 @@ impl Interpreter {
         })
     }
 
+    /// Return a new [`Interpreter`] that should be treated as relocatable.
+    #[must_use]
+    pub fn with_relocatable(self) -> Self {
+        Self {
+            relocatable: true,
+            ..self
+        }
+    }
+
     /// Return the [`Interpreter`] for the base executable, if it's available.
     ///
     /// If no such base executable is available, or if the base executable is the same as the
@@ -460,6 +472,7 @@ impl Interpreter {
                     },
                 }
             },
+            relocatable: self.relocatable,
         }
     }