From cd9e6efe11924f943ea86fdefee9244cdc862c08 Mon Sep 17 00:00:00 2001
From: Pavel Dikov <pavel.dikov@morganstanley.com>
Date: Sat, 27 Jul 2024 19:42:38 +0100
Subject: [PATCH] feat(venv): add relocatable flag (todo: amend activation
 scripts)

---
 crates/install-wheel-rs/src/linker.rs         | 20 ++++-
 crates/install-wheel-rs/src/wheel.rs          | 84 +++++++++++++++----
 crates/uv-build/src/lib.rs                    |  1 +
 crates/uv-cli/src/lib.rs                      | 11 +++
 crates/uv-installer/src/installer.rs          | 13 ++-
 crates/uv-python/src/virtualenv.rs            | 19 ++++-
 crates/uv-tool/src/lib.rs                     |  1 +
 crates/uv-virtualenv/src/lib.rs               |  2 +
 crates/uv-virtualenv/src/virtualenv.rs        |  9 ++
 crates/uv/src/commands/project/environment.rs |  1 +
 crates/uv/src/commands/project/mod.rs         |  1 +
 crates/uv/src/commands/project/run.rs         |  1 +
 crates/uv/src/commands/venv.rs                |  4 +
 crates/uv/src/lib.rs                          |  1 +
 crates/uv/src/settings.rs                     |  3 +
 15 files changed, 152 insertions(+), 19 deletions(-)

diff --git a/crates/install-wheel-rs/src/linker.rs b/crates/install-wheel-rs/src/linker.rs
index 5eb48195ce219..6b00489c4544f 100644
--- a/crates/install-wheel-rs/src/linker.rs
+++ b/crates/install-wheel-rs/src/linker.rs
@@ -46,6 +46,7 @@ pub fn install_wheel(
     installer: Option<&str>,
     link_mode: LinkMode,
     locks: &Locks,
+    is_relocatable: bool,
 ) -> Result<(), Error> {
     let dist_info_prefix = find_dist_info(&wheel)?;
     let metadata = dist_info_metadata(&dist_info_prefix, &wheel)?;
@@ -101,8 +102,22 @@ pub fn install_wheel(
         debug!(name, "Writing entrypoints");
 
         fs_err::create_dir_all(&layout.scheme.scripts)?;
-        write_script_entrypoints(layout, site_packages, &console_scripts, &mut record, false)?;
-        write_script_entrypoints(layout, site_packages, &gui_scripts, &mut record, true)?;
+        write_script_entrypoints(
+            layout,
+            site_packages,
+            &console_scripts,
+            &mut record,
+            false,
+            is_relocatable,
+        )?;
+        write_script_entrypoints(
+            layout,
+            site_packages,
+            &gui_scripts,
+            &mut record,
+            true,
+            is_relocatable,
+        )?;
     }
 
     // 2.a Unpacked archive includes distribution-1.0.dist-info/ and (if there is data) distribution-1.0.data/.
@@ -118,6 +133,7 @@ pub fn install_wheel(
             &console_scripts,
             &gui_scripts,
             &mut record,
+            is_relocatable,
         )?;
         // 2.c If applicable, update scripts starting with #!python to point to the correct interpreter.
         // Script are unsupported through data
diff --git a/crates/install-wheel-rs/src/wheel.rs b/crates/install-wheel-rs/src/wheel.rs
index f53bf8f6ec3fc..149fd0ea410bd 100644
--- a/crates/install-wheel-rs/src/wheel.rs
+++ b/crates/install-wheel-rs/src/wheel.rs
@@ -1,5 +1,5 @@
 use std::collections::HashMap;
-use std::io::{BufRead, BufReader, Cursor, Read, Write};
+use std::io::{BufRead, BufReader, Cursor, Read, Seek, Write};
 use std::path::{Path, PathBuf};
 use std::{env, io};
 
@@ -131,6 +131,7 @@ fn copy_and_hash(reader: &mut impl Read, writer: &mut impl Write) -> io::Result<
 fn format_shebang(executable: impl AsRef<Path>, os_name: &str) -> String {
     // Convert the executable to a simplified path.
     let executable = executable.as_ref().simplified_display().to_string();
+    let is_relocatable = !executable.contains(['\\', '/']);
 
     // Validate the shebang.
     if os_name == "posix" {
@@ -139,11 +140,18 @@ fn format_shebang(executable: impl AsRef<Path>, os_name: &str) -> String {
         let shebang_length = 2 + executable.len() + 1;
 
         // If the shebang is too long, or contains spaces, wrap it in `/bin/sh`.
-        if shebang_length > 127 || executable.contains(' ') {
+        // Same applies for relocatable scripts (executable is relative to script dir, hence `dirname` trick)
+        // (note: the Windows trampoline binaries natively support relative paths to executable)
+        if shebang_length > 127 || executable.contains(' ') || is_relocatable {
+            let prefix = if is_relocatable {
+                r#""$(dirname "$0")/""#
+            } else {
+                ""
+            };
             // Like Python's `shlex.quote`:
             // > Use single quotes, and put single quotes into double quotes
             // > The string $'b is then quoted as '$'"'"'b'
-            let executable = format!("'{}'", executable.replace('\'', r#"'"'"'"#));
+            let executable = format!("{}'{}'", prefix, executable.replace('\'', r#"'"'"'"#));
             return format!("#!/bin/sh\n'''exec' {executable} \"$0\" \"$@\"\n' '''");
         }
     }
@@ -235,9 +243,9 @@ pub(crate) fn windows_script_launcher(
 /// Returns a [`PathBuf`] to `python[w].exe` for script execution.
 ///
 /// <https://github.com/pypa/pip/blob/76e82a43f8fb04695e834810df64f2d9a2ff6020/src/pip/_vendor/distlib/scripts.py#L121-L126>
-fn get_script_executable(python_executable: &Path, is_gui: bool) -> PathBuf {
+fn get_script_executable(python_executable: &Path, is_gui: bool, is_relocatable: bool) -> PathBuf {
     // Only check for pythonw.exe on Windows
-    if cfg!(windows) && is_gui {
+    let script_executable = if cfg!(windows) && is_gui {
         python_executable
             .file_name()
             .map(|name| {
@@ -248,6 +256,14 @@ fn get_script_executable(python_executable: &Path, is_gui: bool) -> PathBuf {
             .unwrap_or_else(|| python_executable.to_path_buf())
     } else {
         python_executable.to_path_buf()
+    };
+    if is_relocatable {
+        script_executable
+            .file_name()
+            .map(PathBuf::from)
+            .unwrap_or_else(|| script_executable)
+    } else {
+        script_executable
     }
 }
 
@@ -276,6 +292,7 @@ pub(crate) fn write_script_entrypoints(
     entrypoints: &[Script],
     record: &mut Vec<RecordEntry>,
     is_gui: bool,
+    is_relocatable: bool,
 ) -> Result<(), Error> {
     for entrypoint in entrypoints {
         let entrypoint_absolute = entrypoint_path(entrypoint, layout);
@@ -292,7 +309,8 @@ pub(crate) fn write_script_entrypoints(
             })?;
 
         // Generate the launcher script.
-        let launcher_executable = get_script_executable(&layout.sys_executable, is_gui);
+        let launcher_executable =
+            get_script_executable(&layout.sys_executable, is_gui, is_relocatable);
         let launcher_python_script = get_script_launcher(
             entrypoint,
             &format_shebang(&launcher_executable, &layout.os_name),
@@ -440,6 +458,7 @@ fn install_script(
     site_packages: &Path,
     record: &mut [RecordEntry],
     file: &DirEntry,
+    is_relocatable: bool,
 ) -> Result<(), Error> {
     let file_type = file.file_type()?;
 
@@ -494,7 +513,18 @@ 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 is_gui = {
+            let mut buf = vec![0];
+            script.read_exact(&mut buf)?;
+            if buf == b"w" {
+                true
+            } else {
+                script.seek_relative(-1)?;
+                false
+            }
+        };
+        let executable = get_script_executable(&layout.sys_executable, is_gui, is_relocatable);
+        let start = format_shebang(&executable, &layout.os_name)
             .as_bytes()
             .to_vec();
 
@@ -561,6 +591,7 @@ pub(crate) fn install_data(
     console_scripts: &[Script],
     gui_scripts: &[Script],
     record: &mut [RecordEntry],
+    is_relocatable: bool,
 ) -> Result<(), Error> {
     for entry in fs::read_dir(data_dir)? {
         let entry = entry?;
@@ -598,7 +629,7 @@ pub(crate) fn install_data(
                         initialized = true;
                     }
 
-                    install_script(layout, site_packages, record, &file)?;
+                    install_script(layout, site_packages, record, &file, is_relocatable)?;
                 }
             }
             Some("headers") => {
@@ -898,6 +929,15 @@ mod test {
             "#!/bin/sh\n'''exec' '/usr/bin/path to python3' \"$0\" \"$@\"\n' '''"
         );
 
+        // And if executable is a relative path, we want a relocatable
+        // script, hence we should use `exec` trick with `dirname`.
+        let executable = Path::new("python3");
+        let os_name = "posix";
+        assert_eq!(
+            format_shebang(executable, os_name),
+            "#!/bin/sh\n'''exec' \"$(dirname \"$0\")/\"'python3' \"$0\" \"$@\"\n' '''"
+        );
+
         // Except on Windows...
         let executable = Path::new("/usr/bin/path to python3");
         let os_name = "nt";
@@ -1005,13 +1045,13 @@ mod test {
         python_exe.write_str("")?;
         pythonw_exe.write_str("")?;
 
-        let script_path = get_script_executable(&python_exe, true);
+        let script_path = get_script_executable(&python_exe, true, false);
         #[cfg(windows)]
         assert_eq!(script_path, pythonw_exe.to_path_buf());
         #[cfg(not(windows))]
         assert_eq!(script_path, python_exe.to_path_buf());
 
-        let script_path = get_script_executable(&python_exe, false);
+        let script_path = get_script_executable(&python_exe, false, false);
         assert_eq!(script_path, python_exe.to_path_buf());
 
         // Test without adjacent pythonw.exe
@@ -1019,10 +1059,10 @@ mod test {
         let python_exe = temp_dir.child("python.exe");
         python_exe.write_str("")?;
 
-        let script_path = get_script_executable(&python_exe, true);
+        let script_path = get_script_executable(&python_exe, true, false);
         assert_eq!(script_path, python_exe.to_path_buf());
 
-        let script_path = get_script_executable(&python_exe, false);
+        let script_path = get_script_executable(&python_exe, false, false);
         assert_eq!(script_path, python_exe.to_path_buf());
 
         // Test with overridden python.exe and pythonw.exe
@@ -1036,15 +1076,31 @@ mod test {
         dot_python_exe.write_str("")?;
         dot_pythonw_exe.write_str("")?;
 
-        let script_path = get_script_executable(&dot_python_exe, true);
+        let script_path = get_script_executable(&dot_python_exe, true, false);
         #[cfg(windows)]
         assert_eq!(script_path, dot_pythonw_exe.to_path_buf());
         #[cfg(not(windows))]
         assert_eq!(script_path, dot_python_exe.to_path_buf());
 
-        let script_path = get_script_executable(&dot_python_exe, false);
+        let script_path = get_script_executable(&dot_python_exe, false, false);
         assert_eq!(script_path, dot_python_exe.to_path_buf());
 
+        // Test with relocatable executable.
+        let temp_dir = assert_fs::TempDir::new()?;
+        let python_exe = temp_dir.child("python.exe");
+        let pythonw_exe = temp_dir.child("pythonw.exe");
+        python_exe.write_str("")?;
+        pythonw_exe.write_str("")?;
+
+        let script_path = get_script_executable(&python_exe, true, true);
+        #[cfg(windows)]
+        assert_eq!(script_path, Path::new("pythonw.exe").to_path_buf());
+        #[cfg(not(windows))]
+        assert_eq!(script_path, Path::new("python.exe").to_path_buf());
+
+        let script_path = get_script_executable(&python_exe, false, true);
+        assert_eq!(script_path, Path::new("python.exe").to_path_buf());
+
         Ok(())
     }
 }
diff --git a/crates/uv-build/src/lib.rs b/crates/uv-build/src/lib.rs
index d629b19c0ea7d..fe84509388a76 100644
--- a/crates/uv-build/src/lib.rs
+++ b/crates/uv-build/src/lib.rs
@@ -442,6 +442,7 @@ impl SourceBuild {
                 uv_virtualenv::Prompt::None,
                 false,
                 false,
+                false,
             )?,
             BuildIsolation::Shared(venv) => venv.clone(),
         };
diff --git a/crates/uv-cli/src/lib.rs b/crates/uv-cli/src/lib.rs
index 27b6809520410..64466f823b62b 100644
--- a/crates/uv-cli/src/lib.rs
+++ b/crates/uv-cli/src/lib.rs
@@ -1743,6 +1743,17 @@ pub struct VenvArgs {
     #[arg(long)]
     pub system_site_packages: bool,
 
+    /// Make the virtual environment relocatable.
+    ///
+    /// A relocatable virtual environment can be moved around and redistributed with its
+    /// associated entrypoint and activation scripts functioning as usual.
+    ///
+    /// Note that this can only be guaranteed for standard `console_scripts` and `gui_scripts`.
+    /// Other scripts may be adjusted if they ship with a generic `#!python[w]` shebang,
+    /// and binaries are left as-is.
+    #[arg(long)]
+    pub relocatable: bool,
+
     #[command(flatten)]
     pub index_args: IndexArgs,
 
diff --git a/crates/uv-installer/src/installer.rs b/crates/uv-installer/src/installer.rs
index 008b8d32543b9..f4ff0d9d9b82d 100644
--- a/crates/uv-installer/src/installer.rs
+++ b/crates/uv-installer/src/installer.rs
@@ -63,9 +63,17 @@ impl<'a> Installer<'a> {
             installer_name,
         } = self;
         let layout = venv.interpreter().layout();
+        let is_relocatable = venv.cfg().is_ok_and(|cfg| cfg.is_relocatable());
 
         rayon::spawn(move || {
-            let result = install(wheels, layout, installer_name, link_mode, reporter);
+            let result = install(
+                wheels,
+                layout,
+                installer_name,
+                link_mode,
+                reporter,
+                is_relocatable,
+            );
             tx.send(result).unwrap();
         });
 
@@ -83,6 +91,7 @@ impl<'a> Installer<'a> {
             self.installer_name,
             self.link_mode,
             self.reporter,
+            self.venv.cfg().is_ok_and(|cfg| cfg.is_relocatable()),
         )
     }
 }
@@ -95,6 +104,7 @@ fn install(
     installer_name: Option<String>,
     link_mode: LinkMode,
     reporter: Option<Box<dyn Reporter>>,
+    is_relocatable: bool,
 ) -> Result<Vec<CachedDist>> {
     let locks = install_wheel_rs::linker::Locks::default();
     wheels.par_iter().try_for_each(|wheel| {
@@ -111,6 +121,7 @@ fn install(
             installer_name.as_deref(),
             link_mode,
             &locks,
+            is_relocatable,
         )
         .with_context(|| format!("Failed to install: {} ({wheel})", wheel.filename()))?;
 
diff --git a/crates/uv-python/src/virtualenv.rs b/crates/uv-python/src/virtualenv.rs
index ed9a862a7cd05..d18265f8ae0bb 100644
--- a/crates/uv-python/src/virtualenv.rs
+++ b/crates/uv-python/src/virtualenv.rs
@@ -28,6 +28,8 @@ pub struct PyVenvConfiguration {
     pub(crate) virtualenv: bool,
     /// If the uv package was used to create the virtual environment.
     pub(crate) uv: bool,
+    /// Is the virtual environment relocatable?
+    pub(crate) relocatable: bool,
 }
 
 #[derive(Debug, Error)]
@@ -136,6 +138,7 @@ impl PyVenvConfiguration {
     pub fn parse(cfg: impl AsRef<Path>) -> Result<Self, Error> {
         let mut virtualenv = false;
         let mut uv = false;
+        let mut relocatable = false;
 
         // Per https://snarky.ca/how-virtual-environments-work/, the `pyvenv.cfg` file is not a
         // valid INI file, and is instead expected to be parsed by partitioning each line on the
@@ -143,7 +146,7 @@ impl PyVenvConfiguration {
         let content = fs::read_to_string(&cfg)
             .map_err(|err| Error::ParsePyVenvCfg(cfg.as_ref().to_path_buf(), err))?;
         for line in content.lines() {
-            let Some((key, _value)) = line.split_once('=') else {
+            let Some((key, value)) = line.split_once('=') else {
                 continue;
             };
             match key.trim() {
@@ -153,11 +156,18 @@ impl PyVenvConfiguration {
                 "uv" => {
                     uv = true;
                 }
+                "relocatable" => {
+                    relocatable = value.trim().to_lowercase() == "true";
+                }
                 _ => {}
             }
         }
 
-        Ok(Self { virtualenv, uv })
+        Ok(Self {
+            virtualenv,
+            uv,
+            relocatable,
+        })
     }
 
     /// Returns true if the virtual environment was created with the `virtualenv` package.
@@ -169,4 +179,9 @@ impl PyVenvConfiguration {
     pub fn is_uv(&self) -> bool {
         self.uv
     }
+
+    /// Returns true if the virtual environment is relocatable.
+    pub fn is_relocatable(&self) -> bool {
+        self.relocatable
+    }
 }
diff --git a/crates/uv-tool/src/lib.rs b/crates/uv-tool/src/lib.rs
index b54bb06f3e41d..c330fd3bd4221 100644
--- a/crates/uv-tool/src/lib.rs
+++ b/crates/uv-tool/src/lib.rs
@@ -258,6 +258,7 @@ impl InstalledTools {
             uv_virtualenv::Prompt::None,
             false,
             false,
+            false,
         )?;
 
         Ok(venv)
diff --git a/crates/uv-virtualenv/src/lib.rs b/crates/uv-virtualenv/src/lib.rs
index 6983284d68149..30e8ed4de9ebe 100644
--- a/crates/uv-virtualenv/src/lib.rs
+++ b/crates/uv-virtualenv/src/lib.rs
@@ -52,6 +52,7 @@ pub fn create_venv(
     prompt: Prompt,
     system_site_packages: bool,
     allow_existing: bool,
+    relocatable: bool,
 ) -> Result<PythonEnvironment, Error> {
     // Create the virtualenv at the given location.
     let virtualenv = virtualenv::create(
@@ -60,6 +61,7 @@ pub fn create_venv(
         prompt,
         system_site_packages,
         allow_existing,
+        relocatable,
     )?;
 
     // Create the corresponding `PythonEnvironment`.
diff --git a/crates/uv-virtualenv/src/virtualenv.rs b/crates/uv-virtualenv/src/virtualenv.rs
index 09348e8e52c12..9374e04360760 100644
--- a/crates/uv-virtualenv/src/virtualenv.rs
+++ b/crates/uv-virtualenv/src/virtualenv.rs
@@ -50,6 +50,7 @@ pub(crate) fn create(
     prompt: Prompt,
     system_site_packages: bool,
     allow_existing: bool,
+    relocatable: bool,
 ) -> Result<VirtualEnvironment, Error> {
     // Determine the base Python executable; that is, the Python executable that should be
     // considered the "base" for the virtual environment. This is typically the Python executable
@@ -335,6 +336,14 @@ pub(crate) fn create(
                 "false".to_string()
             },
         ),
+        (
+            "relocatable".to_string(),
+            if relocatable {
+                "true".to_string()
+            } else {
+                "false".to_string()
+            },
+        ),
     ];
 
     if let Some(prompt) = prompt {
diff --git a/crates/uv/src/commands/project/environment.rs b/crates/uv/src/commands/project/environment.rs
index d6789c5dcb8bb..484bea8d04d8c 100644
--- a/crates/uv/src/commands/project/environment.rs
+++ b/crates/uv/src/commands/project/environment.rs
@@ -132,6 +132,7 @@ impl CachedEnvironment {
             uv_virtualenv::Prompt::None,
             false,
             false,
+            false,
         )?;
 
         let venv = sync_environment(
diff --git a/crates/uv/src/commands/project/mod.rs b/crates/uv/src/commands/project/mod.rs
index 300aa7c80b761..8ba42c0f67093 100644
--- a/crates/uv/src/commands/project/mod.rs
+++ b/crates/uv/src/commands/project/mod.rs
@@ -313,6 +313,7 @@ pub(crate) async fn get_or_init_environment(
                 uv_virtualenv::Prompt::None,
                 false,
                 false,
+                false,
             )?)
         }
     }
diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs
index 76117e368eb13..92b9bf4b811a2 100644
--- a/crates/uv/src/commands/project/run.rs
+++ b/crates/uv/src/commands/project/run.rs
@@ -385,6 +385,7 @@ pub(crate) async fn run(
             uv_virtualenv::Prompt::None,
             false,
             false,
+            false,
         )?;
 
         match spec {
diff --git a/crates/uv/src/commands/venv.rs b/crates/uv/src/commands/venv.rs
index b4abb9f1645d3..c3bb873659abf 100644
--- a/crates/uv/src/commands/venv.rs
+++ b/crates/uv/src/commands/venv.rs
@@ -54,6 +54,7 @@ pub(crate) async fn venv(
     preview: PreviewMode,
     cache: &Cache,
     printer: Printer,
+    relocatable: bool,
 ) -> Result<ExitStatus> {
     match venv_impl(
         path,
@@ -74,6 +75,7 @@ pub(crate) async fn venv(
         native_tls,
         cache,
         printer,
+        relocatable,
     )
     .await
     {
@@ -125,6 +127,7 @@ async fn venv_impl(
     native_tls: bool,
     cache: &Cache,
     printer: Printer,
+    relocatable: bool,
 ) -> miette::Result<ExitStatus> {
     let client_builder = BaseClientBuilder::default()
         .connectivity(connectivity)
@@ -192,6 +195,7 @@ async fn venv_impl(
         prompt,
         system_site_packages,
         allow_existing,
+        relocatable,
     )
     .map_err(VenvError::Creation)?;
 
diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs
index c5ca7eef39569..6a9dbf82e9a90 100644
--- a/crates/uv/src/lib.rs
+++ b/crates/uv/src/lib.rs
@@ -594,6 +594,7 @@ async fn run(cli: Cli) -> Result<ExitStatus> {
                 globals.preview,
                 &cache,
                 printer,
+                args.relocatable,
             )
             .await
         }
diff --git a/crates/uv/src/settings.rs b/crates/uv/src/settings.rs
index cf4e7f5102870..72d392de33788 100644
--- a/crates/uv/src/settings.rs
+++ b/crates/uv/src/settings.rs
@@ -1403,6 +1403,7 @@ pub(crate) struct VenvSettings {
     pub(crate) name: PathBuf,
     pub(crate) prompt: Option<String>,
     pub(crate) system_site_packages: bool,
+    pub(crate) relocatable: bool,
     pub(crate) settings: PipSettings,
 }
 
@@ -1418,6 +1419,7 @@ impl VenvSettings {
             name,
             prompt,
             system_site_packages,
+            relocatable,
             index_args,
             index_strategy,
             keyring_provider,
@@ -1432,6 +1434,7 @@ impl VenvSettings {
             name,
             prompt,
             system_site_packages,
+            relocatable,
             settings: PipSettings::combine(
                 PipOptions {
                     python,