Skip to content

Commit

Permalink
feat(venv): add relocatable flag (todo: amend activation scripts)
Browse files Browse the repository at this point in the history
  • Loading branch information
paveldikov committed Jul 28, 2024
1 parent 866d844 commit cd9e6ef
Show file tree
Hide file tree
Showing 15 changed files with 152 additions and 19 deletions.
20 changes: 18 additions & 2 deletions crates/install-wheel-rs/src/linker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)?;
Expand Down Expand Up @@ -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/.
Expand All @@ -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
Expand Down
84 changes: 70 additions & 14 deletions crates/install-wheel-rs/src/wheel.rs
Original file line number Diff line number Diff line change
@@ -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};

Expand Down Expand Up @@ -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" {
Expand All @@ -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' '''");
}
}
Expand Down Expand Up @@ -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| {
Expand All @@ -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
}
}

Expand Down Expand Up @@ -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);
Expand All @@ -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),
Expand Down Expand Up @@ -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()?;

Expand Down Expand Up @@ -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();

Expand Down Expand Up @@ -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?;
Expand Down Expand Up @@ -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") => {
Expand Down Expand Up @@ -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";
Expand Down Expand Up @@ -1005,24 +1045,24 @@ 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
let temp_dir = assert_fs::TempDir::new()?;
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
Expand All @@ -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(())
}
}
1 change: 1 addition & 0 deletions crates/uv-build/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -442,6 +442,7 @@ impl SourceBuild {
uv_virtualenv::Prompt::None,
false,
false,
false,
)?,
BuildIsolation::Shared(venv) => venv.clone(),
};
Expand Down
11 changes: 11 additions & 0 deletions crates/uv-cli/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,

Expand Down
13 changes: 12 additions & 1 deletion crates/uv-installer/src/installer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
});

Expand All @@ -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()),
)
}
}
Expand All @@ -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| {
Expand All @@ -111,6 +121,7 @@ fn install(
installer_name.as_deref(),
link_mode,
&locks,
is_relocatable,
)
.with_context(|| format!("Failed to install: {} ({wheel})", wheel.filename()))?;

Expand Down
19 changes: 17 additions & 2 deletions crates/uv-python/src/virtualenv.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down Expand Up @@ -136,14 +138,15 @@ 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
// first equals sign.
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() {
Expand All @@ -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.
Expand All @@ -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
}
}
1 change: 1 addition & 0 deletions crates/uv-tool/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -258,6 +258,7 @@ impl InstalledTools {
uv_virtualenv::Prompt::None,
false,
false,
false,
)?;

Ok(venv)
Expand Down
2 changes: 2 additions & 0 deletions crates/uv-virtualenv/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -60,6 +61,7 @@ pub fn create_venv(
prompt,
system_site_packages,
allow_existing,
relocatable,
)?;

// Create the corresponding `PythonEnvironment`.
Expand Down
Loading

0 comments on commit cd9e6ef

Please sign in to comment.