Skip to content

Commit

Permalink
Lookup Python executable from binary
Browse files Browse the repository at this point in the history
  • Loading branch information
MichaReiser committed Feb 21, 2024
1 parent 1b1319b commit ff62bd7
Show file tree
Hide file tree
Showing 8 changed files with 204 additions and 79 deletions.
36 changes: 20 additions & 16 deletions crates/install-wheel-rs/src/wheel.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,13 +33,16 @@ use crate::{find_dist_info, Error};
/// `#!/usr/bin/env python`
pub const SHEBANG_PYTHON: &str = "#!/usr/bin/env python";

const LAUNCHER_MAGIC_NUMBER: [u8; 4] = [b'U', b'V', b'U', b'V'];

#[cfg(all(windows, target_arch = "x86_64"))]
const LAUNCHER_X86_64_GUI: &[u8] =
include_bytes!("../../uv-trampoline/trampolines/uv-trampoline-x86_64-gui.exe");

#[cfg(all(windows, target_arch = "x86_64"))]
const LAUNCHER_X86_64_CONSOLE: &[u8] =
include_bytes!("../../uv-trampoline/trampolines/uv-trampoline-x86_64-console.exe");
const LAUNCHER_X86_64_CONSOLE: &[u8] = include_bytes!(
"../../uv-trampoline/target/x86_64-pc-windows-msvc/release/uv-trampoline-console.exe"
);

#[cfg(all(windows, target_arch = "aarch64"))]
const LAUNCHER_AARCH64_GUI: &[u8] =
Expand Down Expand Up @@ -281,18 +284,7 @@ fn unpack_wheel_files<R: Read + Seek>(
}

fn get_shebang(location: &InstallLocation<impl AsRef<Path>>) -> String {
let path = location.python().to_string_lossy().to_string();
let path = if cfg!(windows) {
// https://stackoverflow.com/a/50323079
const VERBATIM_PREFIX: &str = r"\\?\";
if let Some(stripped) = path.strip_prefix(VERBATIM_PREFIX) {
stripped.to_string()
} else {
path
}
} else {
path
};
let path = location.python().normalized().to_string_lossy().to_string();
format!("#!{path}")
}

Expand All @@ -305,6 +297,7 @@ fn get_shebang(location: &InstallLocation<impl AsRef<Path>>) -> String {
pub(crate) fn windows_script_launcher(
launcher_python_script: &str,
is_gui: bool,
installation: &InstallLocation<impl AsRef<Path>>,
) -> Result<Vec<u8>, Error> {
// This method should only be called on Windows, but we avoid `#[cfg(windows)]` to retain
// compilation on all platforms.
Expand Down Expand Up @@ -352,9 +345,20 @@ pub(crate) fn windows_script_launcher(
archive.finish().expect(error_msg);
}

let python = installation.python();
let python_path = python.normalized().to_string_lossy();

let mut launcher: Vec<u8> = Vec::with_capacity(launcher_bin.len() + payload.len());
launcher.extend_from_slice(launcher_bin);
launcher.extend_from_slice(&payload);
launcher.extend_from_slice(python_path.as_bytes());
launcher.extend_from_slice(
&u32::try_from(python_path.as_bytes().len())
.expect("File Path to be smaller than 4GB")
.to_le_bytes(),
);
launcher.extend_from_slice(&LAUNCHER_MAGIC_NUMBER);

Ok(launcher)
}

Expand Down Expand Up @@ -393,7 +397,7 @@ pub(crate) fn write_script_entrypoints(
write_file_recorded(
site_packages,
&entrypoint_relative,
&windows_script_launcher(&launcher_python_script, is_gui)?,
&windows_script_launcher(&launcher_python_script, is_gui, location)?,
record,
)?;
} else {
Expand Down Expand Up @@ -949,7 +953,7 @@ pub fn parse_key_value_file(
///
/// Wheel 1.0: <https://www.python.org/dev/peps/pep-0427/>
#[allow(clippy::too_many_arguments)]
#[instrument(skip_all, fields(name = %filename.name))]
#[instrument(skip_all, fields(name = % filename.name))]
pub fn install_wheel(
location: &InstallLocation<LockedDir>,
reader: impl Read + Seek,
Expand Down
20 changes: 10 additions & 10 deletions crates/uv-trampoline/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

27 changes: 12 additions & 15 deletions crates/uv-trampoline/README.md
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# Windows trampolines

This is a fork of [posy trampolines](https://github.com/njsmith/posy/tree/dda22e6f90f5fefa339b869dd2bbe107f5b48448/src/trampolines/windows-trampolines/posy-trampoline).
This is a fork
of [posy trampolines](https://github.com/njsmith/posy/tree/dda22e6f90f5fefa339b869dd2bbe107f5b48448/src/trampolines/windows-trampolines/posy-trampoline).

# What is this?

Expand All @@ -13,27 +14,26 @@ That's what this does: it's a generic "trampoline" that lets us generate custom
`.exe`s for arbitrary Python scripts, and when invoked it bounces to invoking
`python <the script>` instead.


# How do you use it?

Basically, this looks up `python.exe` (for console programs) or
`pythonw.exe` (for GUI programs) in the adjacent directory, and invokes
`python[w].exe path\to\the\<the .exe>`.

The intended use is: take your Python script, name it `__main__.py`, and pack it
into a `.zip` file. Then concatenate that `.zip` file onto the end of one of our
prebuilt `.exe`s.
The intended use is:

* take your Python script, name it `__main__.py`, and pack it
into a `.zip` file. Then concatenate that `.zip` file onto the end of one of our
prebuilt `.exe`s.
* After the zip file content, write the path to the Python executable that the script uses to run
the Python script as UTF-8 encoded string, followed by the path's length as a 32-bit little-endian
integer.
* At the very end, write the magic number `UVUV` in bytes.

Then when you run `python` on the `.exe`, it will see the `.zip` trailer at the
end of the `.exe`, and automagically look inside to find and execute
`__main__.py`. Easy-peasy.

(TODO: we should probably make the Python-finding logic slightly more flexible
at some point -- in particular to support more conventional venv-style
installation where you find `python` by looking in the directory next to the
trampoline `.exe` -- but this is good enough to get started.)


# Why does this exist?

I probably could have used Vinay's C++ implementation from `distlib`, but what's
Expand All @@ -47,7 +47,6 @@ Python-finding logic we want. But mostly it was just an interesting challenge.
This does owe a *lot* to the `distlib` implementation though. The overall logic
is copied more-or-less directly.


# Anything I should know for hacking on this?

In order to minimize binary size, this uses `#![no_std]`, `panic="abort"`, and
Expand All @@ -64,7 +63,7 @@ this:
Though uh, this does mean that literally all of our code is `unsafe`. Sorry!

- `runtime.rs` has the core glue to get panicking, heap allocation, and linking
working.
working.

- `diagnostics.rs` uses `ufmt` and some cute Windows tricks to get a convenient
version of `eprintln!` that works without `std`, and automatically prints to
Expand All @@ -85,7 +84,6 @@ Miscellaneous tips:
`.unwrap_unchecked()` avoids this. Similar for `slice[idx]` vs
`slice.get_unchecked(idx)`.


# How do you build this stupid thing?

Building this can be frustrating, because the low-level compiler/runtime
Expand All @@ -107,7 +105,6 @@ Two approaches that are reasonably likely to work:
- Leave `compiler-builtins` commented-out, and build like: `cargo build
--release -Z build-std=core,panic_abort,alloc -Z
build-std-features=compiler-builtins-mem --target x86_64-pc-windows-msvc`


Hopefully in the future as `#![no_std]` develops, this will get smoother.

Expand Down
Loading

0 comments on commit ff62bd7

Please sign in to comment.