Skip to content

Commit 34a4330

Browse files
zaniebzanieb
committed
Add a trampoline variant that just executes python (#8637)
Currently, our trampoline is used to convert `<command> [args]` to `python <command> [args]` for script entrypoints installed into virtual environments. For #8458, it'd be nice to convert a shim `python3.12 [args]` to `python [args]`. Here, we modify the trampolines to support this use-case. The only change we really need here is to avoid injecting `<command>` into the child process. We change the "magic number" at the end of the trampoline executables from `UVUV` to `UVSC` and `UVPY` which define "script" and "python" variants to the trampoline. We then omit the `<command>` injection in the latter case. We also omit writing the zip script payload. To support construction of the new variant, a new `uv-trampoline-builder` crate is introduced — this avoids requirements on `uv-install-wheel` in future work. I also use `uv-trampoline-builder` to consolidate some of the test setup for `uv-trampoline`. There should be no backwards compatibility concerns, since trampolines are fully self-referential. I rebased to fix the commits at the end, as this took many iterations to get working via CI. This should roughly be reviewable by commit if you prefer.
1 parent 4579f3f commit 34a4330

19 files changed

+564
-474
lines changed

‎.github/workflows/ci.yml

+10-3
Original file line numberDiff line numberDiff line change
@@ -396,16 +396,23 @@ jobs:
396396
- uses: Swatinem/rust-cache@v2
397397
with:
398398
workspaces: ${{ github.workspace }}/crates/uv-trampoline
399+
- name: "Test committed binaries"
400+
working-directory: ${{ github.workspace }}
401+
run: |
402+
rustup target add ${{ matrix.target-arch }}-pc-windows-msvc
403+
cargo test -p uv-trampoline-builder --target ${{ matrix.target-arch }}-pc-windows-msvc
399404
# Build and copy the new binaries
400405
- name: "Build"
401406
working-directory: ${{ github.workspace }}/crates/uv-trampoline
402407
run: |
403408
cargo build --target ${{ matrix.target-arch }}-pc-windows-msvc
404409
cp target/${{ matrix.target-arch }}-pc-windows-msvc/debug/uv-trampoline-console.exe trampolines/uv-trampoline-${{ matrix.target-arch }}-console.exe
405410
cp target/${{ matrix.target-arch }}-pc-windows-msvc/debug/uv-trampoline-gui.exe trampolines/uv-trampoline-${{ matrix.target-arch }}-gui.exe
406-
- name: "Test"
407-
working-directory: ${{ github.workspace }}/crates/uv-trampoline
408-
run: cargo test --target ${{ matrix.target-arch }}-pc-windows-msvc --test *
411+
- name: "Test new binaries"
412+
working-directory: ${{ github.workspace }}
413+
run: |
414+
# We turn off the default "production" test feature since these are debug binaries
415+
cargo test -p uv-trampoline-builder --target ${{ matrix.target-arch }}-pc-windows-msvc --no-default-features
409416
410417
typos:
411418
runs-on: ubuntu-latest

‎Cargo.lock

+15
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml

+1
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ uv-settings = { path = "crates/uv-settings" }
6161
uv-shell = { path = "crates/uv-shell" }
6262
uv-state = { path = "crates/uv-state" }
6363
uv-static = { path = "crates/uv-static" }
64+
uv-trampoline-builder = { path = "crates/uv-trampoline-builder" }
6465
uv-tool = { path = "crates/uv-tool" }
6566
uv-types = { path = "crates/uv-types" }
6667
uv-version = { path = "crates/uv-version" }

‎crates/uv-install-wheel/Cargo.toml

+1
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ uv-normalize = { workspace = true }
2828
uv-pep440 = { workspace = true }
2929
uv-platform-tags = { workspace = true }
3030
uv-pypi-types = { workspace = true }
31+
uv-trampoline-builder = { workspace = true }
3132
uv-warnings = { workspace = true }
3233

3334
clap = { workspace = true, optional = true, features = ["derive"] }

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

+2
Original file line numberDiff line numberDiff line change
@@ -97,4 +97,6 @@ pub enum Error {
9797
MismatchedVersion(Version, Version),
9898
#[error("Invalid egg-link")]
9999
InvalidEggLink(PathBuf),
100+
#[error(transparent)]
101+
LauncherError(#[from] uv_trampoline_builder::Error),
100102
}

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

+8-162
Original file line numberDiff line numberDiff line change
@@ -1,51 +1,26 @@
11
use std::collections::HashMap;
2-
use std::io::{BufReader, Cursor, Read, Seek, Write};
2+
use std::io;
3+
use std::io::{BufReader, Read, Seek, Write};
34
use std::path::{Path, PathBuf};
4-
use std::{env, io};
55

6-
use crate::record::RecordEntry;
7-
use crate::script::Script;
8-
use crate::{Error, Layout};
96
use data_encoding::BASE64URL_NOPAD;
107
use fs_err as fs;
118
use fs_err::{DirEntry, File};
129
use mailparse::parse_headers;
1310
use rustc_hash::FxHashMap;
1411
use sha2::{Digest, Sha256};
1512
use tracing::{instrument, warn};
13+
use walkdir::WalkDir;
14+
1615
use uv_cache_info::CacheInfo;
1716
use uv_fs::{relative_to, Simplified};
1817
use uv_normalize::PackageName;
1918
use uv_pypi_types::DirectUrl;
20-
use walkdir::WalkDir;
21-
use zip::write::FileOptions;
22-
use zip::ZipWriter;
23-
24-
const LAUNCHER_MAGIC_NUMBER: [u8; 4] = [b'U', b'V', b'U', b'V'];
19+
use uv_trampoline_builder::windows_script_launcher;
2520

26-
#[cfg(all(windows, target_arch = "x86"))]
27-
const LAUNCHER_I686_GUI: &[u8] =
28-
include_bytes!("../../uv-trampoline/trampolines/uv-trampoline-i686-gui.exe");
29-
30-
#[cfg(all(windows, target_arch = "x86"))]
31-
const LAUNCHER_I686_CONSOLE: &[u8] =
32-
include_bytes!("../../uv-trampoline/trampolines/uv-trampoline-i686-console.exe");
33-
34-
#[cfg(all(windows, target_arch = "x86_64"))]
35-
const LAUNCHER_X86_64_GUI: &[u8] =
36-
include_bytes!("../../uv-trampoline/trampolines/uv-trampoline-x86_64-gui.exe");
37-
38-
#[cfg(all(windows, target_arch = "x86_64"))]
39-
const LAUNCHER_X86_64_CONSOLE: &[u8] =
40-
include_bytes!("../../uv-trampoline/trampolines/uv-trampoline-x86_64-console.exe");
41-
42-
#[cfg(all(windows, target_arch = "aarch64"))]
43-
const LAUNCHER_AARCH64_GUI: &[u8] =
44-
include_bytes!("../../uv-trampoline/trampolines/uv-trampoline-aarch64-gui.exe");
45-
46-
#[cfg(all(windows, target_arch = "aarch64"))]
47-
const LAUNCHER_AARCH64_CONSOLE: &[u8] =
48-
include_bytes!("../../uv-trampoline/trampolines/uv-trampoline-aarch64-console.exe");
21+
use crate::record::RecordEntry;
22+
use crate::script::Script;
23+
use crate::{Error, Layout};
4924

5025
/// Wrapper script template function
5126
///
@@ -158,87 +133,6 @@ fn format_shebang(executable: impl AsRef<Path>, os_name: &str, relocatable: bool
158133
format!("#!{executable}")
159134
}
160135

161-
/// A Windows script is a minimal .exe launcher binary with the python entrypoint script appended as
162-
/// stored zip file.
163-
///
164-
/// <https://github.com/pypa/pip/blob/fd0ea6bc5e8cb95e518c23d901c26ca14db17f89/src/pip/_vendor/distlib/scripts.py#L248-L262>
165-
#[allow(unused_variables)]
166-
pub(crate) fn windows_script_launcher(
167-
launcher_python_script: &str,
168-
is_gui: bool,
169-
python_executable: impl AsRef<Path>,
170-
) -> Result<Vec<u8>, Error> {
171-
// This method should only be called on Windows, but we avoid `#[cfg(windows)]` to retain
172-
// compilation on all platforms.
173-
if cfg!(not(windows)) {
174-
return Err(Error::NotWindows);
175-
}
176-
177-
let launcher_bin: &[u8] = match env::consts::ARCH {
178-
#[cfg(all(windows, target_arch = "x86"))]
179-
"x86" => {
180-
if is_gui {
181-
LAUNCHER_I686_GUI
182-
} else {
183-
LAUNCHER_I686_CONSOLE
184-
}
185-
}
186-
#[cfg(all(windows, target_arch = "x86_64"))]
187-
"x86_64" => {
188-
if is_gui {
189-
LAUNCHER_X86_64_GUI
190-
} else {
191-
LAUNCHER_X86_64_CONSOLE
192-
}
193-
}
194-
#[cfg(all(windows, target_arch = "aarch64"))]
195-
"aarch64" => {
196-
if is_gui {
197-
LAUNCHER_AARCH64_GUI
198-
} else {
199-
LAUNCHER_AARCH64_CONSOLE
200-
}
201-
}
202-
#[cfg(windows)]
203-
arch => {
204-
return Err(Error::UnsupportedWindowsArch(arch));
205-
}
206-
#[cfg(not(windows))]
207-
arch => &[],
208-
};
209-
210-
let mut payload: Vec<u8> = Vec::new();
211-
{
212-
// We're using the zip writer, but with stored compression
213-
// https://github.com/njsmith/posy/blob/04927e657ca97a5e35bb2252d168125de9a3a025/src/trampolines/mod.rs#L75-L82
214-
// https://github.com/pypa/distlib/blob/8ed03aab48add854f377ce392efffb79bb4d6091/PC/launcher.c#L259-L271
215-
let stored = FileOptions::default().compression_method(zip::CompressionMethod::Stored);
216-
let mut archive = ZipWriter::new(Cursor::new(&mut payload));
217-
let error_msg = "Writing to Vec<u8> should never fail";
218-
archive.start_file("__main__.py", stored).expect(error_msg);
219-
archive
220-
.write_all(launcher_python_script.as_bytes())
221-
.expect(error_msg);
222-
archive.finish().expect(error_msg);
223-
}
224-
225-
let python = python_executable.as_ref();
226-
let python_path = python.simplified_display().to_string();
227-
228-
let mut launcher: Vec<u8> = Vec::with_capacity(launcher_bin.len() + payload.len());
229-
launcher.extend_from_slice(launcher_bin);
230-
launcher.extend_from_slice(&payload);
231-
launcher.extend_from_slice(python_path.as_bytes());
232-
launcher.extend_from_slice(
233-
&u32::try_from(python_path.as_bytes().len())
234-
.expect("File Path to be smaller than 4GB")
235-
.to_le_bytes(),
236-
);
237-
launcher.extend_from_slice(&LAUNCHER_MAGIC_NUMBER);
238-
239-
Ok(launcher)
240-
}
241-
242136
/// Returns a [`PathBuf`] to `python[w].exe` for script execution.
243137
///
244138
/// <https://github.com/pypa/pip/blob/76e82a43f8fb04695e834810df64f2d9a2ff6020/src/pip/_vendor/distlib/scripts.py#L121-L126>
@@ -1075,54 +969,6 @@ mod test {
1075969
Ok(())
1076970
}
1077971

1078-
#[test]
1079-
#[cfg(all(windows, target_arch = "x86"))]
1080-
fn test_launchers_are_small() {
1081-
// At time of writing, they are 45kb~ bytes.
1082-
assert!(
1083-
super::LAUNCHER_I686_GUI.len() < 45 * 1024,
1084-
"GUI launcher: {}",
1085-
super::LAUNCHER_I686_GUI.len()
1086-
);
1087-
assert!(
1088-
super::LAUNCHER_I686_CONSOLE.len() < 45 * 1024,
1089-
"CLI launcher: {}",
1090-
super::LAUNCHER_I686_CONSOLE.len()
1091-
);
1092-
}
1093-
1094-
#[test]
1095-
#[cfg(all(windows, target_arch = "x86_64"))]
1096-
fn test_launchers_are_small() {
1097-
// At time of writing, they are 45kb~ bytes.
1098-
assert!(
1099-
super::LAUNCHER_X86_64_GUI.len() < 45 * 1024,
1100-
"GUI launcher: {}",
1101-
super::LAUNCHER_X86_64_GUI.len()
1102-
);
1103-
assert!(
1104-
super::LAUNCHER_X86_64_CONSOLE.len() < 45 * 1024,
1105-
"CLI launcher: {}",
1106-
super::LAUNCHER_X86_64_CONSOLE.len()
1107-
);
1108-
}
1109-
1110-
#[test]
1111-
#[cfg(all(windows, target_arch = "aarch64"))]
1112-
fn test_launchers_are_small() {
1113-
// At time of writing, they are 45kb~ bytes.
1114-
assert!(
1115-
super::LAUNCHER_AARCH64_GUI.len() < 45 * 1024,
1116-
"GUI launcher: {}",
1117-
super::LAUNCHER_AARCH64_GUI.len()
1118-
);
1119-
assert!(
1120-
super::LAUNCHER_AARCH64_CONSOLE.len() < 45 * 1024,
1121-
"CLI launcher: {}",
1122-
super::LAUNCHER_AARCH64_CONSOLE.len()
1123-
);
1124-
}
1125-
1126972
#[test]
1127973
fn test_script_executable() -> Result<()> {
1128974
// Test with adjacent pythonw.exe

‎crates/uv-trampoline-builder/Cargo.toml

+34
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
[package]
2+
name = "uv-trampoline-builder"
3+
version = "0.0.1"
4+
publish = false
5+
description = "Builds launchers for `uv-trampoline`"
6+
7+
edition = { workspace = true }
8+
rust-version = { workspace = true }
9+
homepage = { workspace = true }
10+
documentation = { workspace = true }
11+
repository = { workspace = true }
12+
authors = { workspace = true }
13+
license = { workspace = true }
14+
15+
[features]
16+
default = ["production"]
17+
18+
# Expect tests to run against production builds of `uv-trampoline` binaries, rather than debug builds
19+
production = []
20+
21+
[lints]
22+
workspace = true
23+
24+
[dependencies]
25+
uv-fs = { workspace = true }
26+
thiserror = { workspace = true }
27+
zip = { workspace = true }
28+
29+
[dev-dependencies]
30+
assert_cmd = { version = "2.0.16" }
31+
assert_fs = { version = "1.1.2" }
32+
anyhow = { version = "1.0.89" }
33+
fs-err = { workspace = true }
34+
which = { workspace = true }

0 commit comments

Comments
 (0)