Skip to content

Commit 6bcb8e1

Browse files
charliermarshcharliermarsh
committed
Add relocatable installs to support concurrency-safe cached environments
1 parent ae11317 commit 6bcb8e1

7 files changed

Lines changed: 145 additions & 69 deletions

File tree

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
//! Takes a wheel and installs it into a venv.
22
33
use std::io;
4-
54
use std::path::PathBuf;
65

76
use platform_info::PlatformInfoError;
@@ -34,6 +33,11 @@ pub struct Layout {
3433
pub os_name: String,
3534
/// The [`Scheme`] paths for the interpreter.
3635
pub scheme: Scheme,
36+
/// Whether the environment is "relocatable". When enabled, any paths to the environment (e.g.,
37+
/// those encoded in entrypoints and scripts) will be expressed in relative terms. As a result,
38+
/// the entrypoints and scripts themselves will _not_ be relocatable, but the environment as a
39+
/// whole will be.
40+
pub relocatable: bool,
3741
}
3842

3943
/// Note: The caller is responsible for adding the path of the wheel we're installing.

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

Lines changed: 101 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
1-
use std::collections::HashMap;
2-
use std::io::{BufRead, BufReader, Cursor, Read, Write};
3-
use std::path::{Path, PathBuf};
4-
use std::{env, io};
5-
61
use data_encoding::BASE64URL_NOPAD;
72
use fs_err as fs;
83
use fs_err::{DirEntry, File};
94
use mailparse::MailHeaderMap;
105
use rustc_hash::FxHashMap;
116
use sha2::{Digest, Sha256};
7+
use std::borrow::Cow;
8+
use std::collections::HashMap;
9+
use std::io::{BufRead, BufReader, Cursor, Read, Write};
10+
use std::path::{Path, PathBuf};
11+
use std::{env, io};
1212
use tracing::{instrument, warn};
1313
use walkdir::WalkDir;
1414
use zip::write::FileOptions;
@@ -122,15 +122,22 @@ fn copy_and_hash(reader: &mut impl Read, writer: &mut impl Write) -> io::Result<
122122
))
123123
}
124124

125-
/// Format the shebang for a given Python executable.
125+
/// Format the shebang for a given Python executable, assuming an absolute path.
126126
///
127127
/// Like pip, if a shebang is non-simple (too long or contains spaces), we use `/bin/sh` as the
128128
/// executable.
129129
///
130130
/// See: <https://github.com/pypa/pip/blob/0ad4c94be74cc24874c6feb5bb3c2152c398a18e/src/pip/_vendor/distlib/scripts.py#L136-L165>
131-
fn format_shebang(executable: impl AsRef<Path>, os_name: &str) -> String {
131+
fn format_absolute_shebang(executable: impl AsRef<Path>, os_name: &str) -> String {
132+
let executable = executable.as_ref();
133+
debug_assert!(
134+
executable.is_absolute(),
135+
"Path must be absolute: {}",
136+
executable.display()
137+
);
138+
132139
// Convert the executable to a simplified path.
133-
let executable = executable.as_ref().simplified_display().to_string();
140+
let executable = executable.simplified_display().to_string();
134141

135142
// Validate the shebang.
136143
if os_name == "posix" {
@@ -151,6 +158,32 @@ fn format_shebang(executable: impl AsRef<Path>, os_name: &str) -> String {
151158
format!("#!{executable}")
152159
}
153160

161+
/// Format the shebang for a given Python executable, assuming a relative path.
162+
///
163+
/// Like pip, if a shebang is non-simple (too long or contains spaces), we use `/bin/sh` as the
164+
/// executable.
165+
///
166+
/// See: <https://github.com/pypa/pip/blob/0ad4c94be74cc24874c6feb5bb3c2152c398a18e/src/pip/_vendor/distlib/scripts.py#L136-L165>
167+
fn format_relative_shebang(executable: impl AsRef<Path>, os_name: &str) -> String {
168+
let executable = executable.as_ref();
169+
debug_assert!(
170+
executable.is_relative(),
171+
"Path must be relative: {}",
172+
executable.display()
173+
);
174+
175+
// Convert the executable to a simplified path.
176+
let executable = executable.simplified_display().to_string();
177+
178+
// Wrap in `dirname`. We assume that the relative path is fairly simple, since we know it's a
179+
// relative path within a virtual environment, and so we shouldn't need to handle quotes,e tc.
180+
if os_name == "posix" {
181+
return format!("#!/bin/sh\n'''exec' \"$(dirname $0)/{executable}\" \"$0\" \"$@\"\n' '''");
182+
}
183+
184+
format!("#!{executable}")
185+
}
186+
154187
/// A Windows script is a minimal .exe launcher binary with the python entrypoint script appended as
155188
/// stored zip file.
156189
///
@@ -291,12 +324,32 @@ pub(crate) fn write_script_entrypoints(
291324
))
292325
})?;
293326

327+
let python_executable = if layout.relocatable {
328+
Cow::Owned(
329+
pathdiff::diff_paths(&layout.sys_executable, &layout.scheme.scripts).ok_or_else(
330+
|| {
331+
Error::Io(io::Error::new(
332+
io::ErrorKind::Other,
333+
format!(
334+
"Could not find relative path for: {}",
335+
layout.sys_executable.simplified_display()
336+
),
337+
))
338+
},
339+
)?,
340+
)
341+
} else {
342+
Cow::Borrowed(&layout.sys_executable)
343+
};
344+
294345
// Generate the launcher script.
295-
let launcher_executable = get_script_executable(&layout.sys_executable, is_gui);
296-
let launcher_python_script = get_script_launcher(
297-
entrypoint,
298-
&format_shebang(&launcher_executable, &layout.os_name),
299-
);
346+
let launcher_executable = get_script_executable(&python_executable, is_gui);
347+
let shebang = if layout.relocatable {
348+
format_relative_shebang(&launcher_executable, &layout.os_name)
349+
} else {
350+
format_absolute_shebang(&launcher_executable, &layout.os_name)
351+
};
352+
let launcher_python_script = get_script_launcher(entrypoint, &shebang);
300353

301354
// If necessary, wrap the launcher script in a Windows launcher binary.
302355
if cfg!(windows) {
@@ -432,9 +485,9 @@ pub(crate) fn move_folder_recorded(
432485
Ok(())
433486
}
434487

435-
/// Installs a single script (not an entrypoint)
488+
/// Installs a single script (not an entrypoint).
436489
///
437-
/// Has to deal with both binaries files (just move) and scripts (rewrite the shebang if applicable)
490+
/// Has to deal with both binaries files (just move) and scripts (rewrite the shebang if applicable).
438491
fn install_script(
439492
layout: &Layout,
440493
site_packages: &Path,
@@ -494,7 +547,25 @@ fn install_script(
494547
let mut start = vec![0; placeholder_python.len()];
495548
script.read_exact(&mut start)?;
496549
let size_and_encoded_hash = if start == placeholder_python {
497-
let start = format_shebang(&layout.sys_executable, &layout.os_name)
550+
let python_executable = if layout.relocatable {
551+
Cow::Owned(
552+
pathdiff::diff_paths(&layout.sys_executable, &layout.scheme.scripts).ok_or_else(
553+
|| {
554+
Error::Io(io::Error::new(
555+
io::ErrorKind::Other,
556+
format!(
557+
"Could not find relative path for: {}",
558+
layout.sys_executable.simplified_display()
559+
),
560+
))
561+
},
562+
)?,
563+
)
564+
} else {
565+
Cow::Borrowed(&layout.sys_executable)
566+
};
567+
568+
let start = format_absolute_shebang(python_executable.as_ref(), &layout.os_name)
498569
.as_bytes()
499570
.to_vec();
500571

@@ -779,7 +850,7 @@ mod test {
779850
use assert_fs::prelude::*;
780851
use indoc::{formatdoc, indoc};
781852

782-
use crate::wheel::format_shebang;
853+
use crate::wheel::format_absolute_shebang;
783854
use crate::Error;
784855

785856
use super::{
@@ -884,37 +955,44 @@ mod test {
884955
}
885956

886957
#[test]
887-
fn test_shebang() {
958+
#[cfg(not(windows))]
959+
fn test_absolute() {
888960
// By default, use a simple shebang.
889961
let executable = Path::new("/usr/bin/python3");
890962
let os_name = "posix";
891-
assert_eq!(format_shebang(executable, os_name), "#!/usr/bin/python3");
963+
assert_eq!(
964+
format_absolute_shebang(executable, os_name),
965+
"#!/usr/bin/python3"
966+
);
892967

893968
// If the path contains spaces, we should use the `exec` trick.
894969
let executable = Path::new("/usr/bin/path to python3");
895970
let os_name = "posix";
896971
assert_eq!(
897-
format_shebang(executable, os_name),
972+
format_absolute_shebang(executable, os_name),
898973
"#!/bin/sh\n'''exec' '/usr/bin/path to python3' \"$0\" \"$@\"\n' '''"
899974
);
900975

901976
// Except on Windows...
902977
let executable = Path::new("/usr/bin/path to python3");
903978
let os_name = "nt";
904979
assert_eq!(
905-
format_shebang(executable, os_name),
980+
format_absolute_shebang(executable, os_name),
906981
"#!/usr/bin/path to python3"
907982
);
908983

909984
// Quotes, however, are ok.
910985
let executable = Path::new("/usr/bin/'python3'");
911986
let os_name = "posix";
912-
assert_eq!(format_shebang(executable, os_name), "#!/usr/bin/'python3'");
987+
assert_eq!(
988+
format_absolute_shebang(executable, os_name),
989+
"#!/usr/bin/'python3'"
990+
);
913991

914992
// If the path is too long, we should not use the `exec` trick.
915993
let executable = Path::new("/usr/bin/path/to/a/very/long/executable/executable/executable/executable/executable/executable/executable/executable/name/python3");
916994
let os_name = "posix";
917-
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' '''");
995+
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' '''");
918996
}
919997

920998
#[test]

‎crates/uv-installer/src/installer.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
1-
use std::convert;
2-
31
use anyhow::{Context, Error, Result};
42
use install_wheel_rs::{linker::LinkMode, Layout};
53
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
4+
use std::convert;
65
use tokio::sync::oneshot;
76
use tracing::instrument;
87

‎crates/uv-python/src/environment.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,16 @@ impl PythonEnvironment {
144144
})))
145145
}
146146

147+
/// Create a [`PythonEnvironment`] from an existing [`Interpreter`] with relocatable paths.
148+
#[must_use]
149+
pub fn with_relocatable(self) -> Self {
150+
let inner = Arc::unwrap_or_clone(self.0);
151+
Self(Arc::new(PythonEnvironmentShared {
152+
interpreter: inner.interpreter.with_relocatable(),
153+
..inner
154+
}))
155+
}
156+
147157
/// Returns the root (i.e., `prefix`) of the Python interpreter.
148158
pub fn root(&self) -> &Path {
149159
&self.0.root

‎crates/uv-python/src/interpreter.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ pub struct Interpreter {
4545
prefix: Option<Prefix>,
4646
pointer_size: PointerSize,
4747
gil_disabled: bool,
48+
relocatable: bool,
4849
}
4950

5051
impl Interpreter {
@@ -75,6 +76,7 @@ impl Interpreter {
7576
tags: OnceLock::new(),
7677
target: None,
7778
prefix: None,
79+
relocatable: false,
7880
})
7981
}
8082

@@ -109,6 +111,7 @@ impl Interpreter {
109111
prefix: None,
110112
pointer_size: PointerSize::_64,
111113
gil_disabled: false,
114+
relocatable: false,
112115
}
113116
}
114117

@@ -143,6 +146,15 @@ impl Interpreter {
143146
})
144147
}
145148

149+
/// Return a new [`Interpreter`] that should be treated as relocatable.
150+
#[must_use]
151+
pub fn with_relocatable(self) -> Self {
152+
Self {
153+
relocatable: true,
154+
..self
155+
}
156+
}
157+
146158
/// Return the [`Interpreter`] for the base executable, if it's available.
147159
///
148160
/// If no such base executable is available, or if the base executable is the same as the
@@ -460,6 +472,7 @@ impl Interpreter {
460472
},
461473
}
462474
},
475+
relocatable: self.relocatable,
463476
}
464477
}
465478

0 commit comments

Comments
 (0)