Windows launchers using posy trampolines (astral-sh#1092)

## Background

In virtual environments, we want to install python programs as console
commands, e.g. `black .` over `python -m black .`. They may be called
[entrypoints](https://packaging.python.org/en/latest/specifications/entry-points/)
or scripts. For entrypoints, we're given a module name and function to
call in that module.

On Unix, we generate a minimal python script launcher. Text files are
runnable on unix by adding a shebang at their top, e.g.

```python
#!/usr/bin/env python
```

will make the operating system run the file with the current python
interpreter. A venv launcher for black in `/home/ferris/colorize/.venv`
(module name: `black`, function to call: `patched_main`) would look like
this:

```python
#!/home/ferris/colorize/.venv/bin/python
# -*- coding: utf-8 -*-
import re
import sys
from black import patched_main
if __name__ == "__main__":
    sys.argv[0] = re.sub(r"(-script\.pyw|\.exe)?$", "", sys.argv[0])
    sys.exit(patched_main())
```

On windows, this doesn't work, we can only rely on launching `.exe`
files.

## Summary

We use posy's rust implementation of a trampoline, which is based on
distlib's c++ implementation. We pre-build a minimal exe and append the
launcher script as stored zip archive behind it. The exe will look for
the venv python interpreter next to it and use it to execute the
appended script.

The changes in this PR make the `black` entrypoint work:

```powershell
cargo run -- venv .venv
cargo run -q -- pip install black
.\.venv\Scripts\black --version
```

Integration with our existing tests will be done in follow-up PRs.

## Implementation and Details

I've vendored the posy trampoline crate. It is a formatted, renamed and
slightly changed for embedding version of
njsmith/posy#28.

The posy launchers are smaller than the distlib launchers, 16K vs 106K
for black. Currently only `x86_64-pc-windows-msvc` is supported. The
crate requires a nightly compiler for its no-std binary size tricks.

On windows, an application can be launched with a console or without (to
create windows instead), which needs two different launchers. The gui
launcher will subsequently use `pythonw.exe` while the console launcher
uses `python.exe`.

.github/workflows/ci.yml

@@ -98,3 +98,25 @@ jobs:
       - uses: rui314/setup-mold@v1
       - uses: Swatinem/rust-cache@v2
       - run: cargo clippy --workspace --all-targets --all-features --locked -- -D warnings
+
+  # Separate job for the nightly crate
+  windows-trampoline:
+    runs-on: windows-latest
+    name: "check windows trampoline"
+    steps:
+      - uses: actions/checkout@v4
+      - name: "Install Rust toolchain"
+        working-directory: crates/puffin-trampoline
+        run: |
+          rustup target add x86_64-pc-windows-msvc
+          rustup component add clippy rust-src --toolchain nightly-2024-01-23-x86_64-pc-windows-msvc
+      - uses: rui314/setup-mold@v1
+      - uses: Swatinem/rust-cache@v2
+        with:
+          workspaces: "crates/puffin-trampoline"
+      - name: "Clippy"
+        working-directory: crates/puffin-trampoline
+        run: cargo clippy --all-features --locked -- -D warnings
+      - name: "Build"
+        working-directory: crates/puffin-trampoline
+        run: cargo build --release -Z build-std=core,panic_abort,alloc -Z build-std-features=compiler-builtins-mem --target x86_64-pc-windows-msvc

Cargo.toml

@@ -1,6 +1,10 @@
 [workspace]
 members = ["crates/*"]
-exclude = ["scripts"]
+exclude = [
+    "scripts",
+    # Needs nightly
+    "crates/puffin-trampoline"
+]
 resolver = "2"
 
 [workspace.package]

README.md

@@ -263,7 +263,8 @@ Puffin's Git implementation draws on details from [Cargo](https://github.com/rus
 
 Some of Puffin's optimizations are inspired by the great work we've seen in
 [Orogene](https://github.com/orogene/orogene) and [Bun](https://github.com/oven-sh/bun). We've also
-learned a lot from [Posy](https://github.com/njsmith/posy).
+learned a lot from Nathaniel J. Smith's [Posy](https://github.com/njsmith/posy) and adapted its
+[trampoline](https://github.com/njsmith/posy/tree/main/src/trampolines/windows-trampolines/posy-trampoline).
 
 ## License
 

crates/gourgeist/src/bare.rs

@@ -119,21 +119,26 @@ pub fn create_bare_venv(location: &Utf8Path, interpreter: &Interpreter) -> io::R
     {
         // https://github.com/python/cpython/blob/d457345bbc6414db0443819290b04a9a4333313d/Lib/venv/__init__.py#L261-L267
         // https://github.com/pypa/virtualenv/blob/d9fdf48d69f0d0ca56140cf0381edbb5d6fe09f5/src/virtualenv/create/via_global_ref/builtin/cpython/cpython3.py#L78-L83
-        let shim = interpreter
-            .stdlib()
-            .join("venv")
-            .join("scripts")
-            .join("nt")
-            .join("python.exe");
-        fs_err::copy(shim, bin_dir.join("python.exe"))?;
+        // There's two kinds of applications on windows: Those that allocate a console (python.exe) and those that
+        // don't because they use window(s) (pythonw.exe).
+        for python_exe in ["python.exe", "pythonw.exe"] {
+            let shim = interpreter
+                .stdlib()
+                .join("venv")
+                .join("scripts")
+                .join("nt")
+                .join(python_exe);
+            fs_err::copy(shim, bin_dir.join(python_exe))?;
+        }
     }
     #[cfg(not(any(unix, windows)))]
     {
         compile_error!("Only Windows and Unix are supported")
     }
 
     // Add all the activate scripts for different shells
-    // TODO(konstin): That's unix!
+    // TODO(konstin): RELATIVE_SITE_PACKAGES is currently only the unix path. We should ensure that all launchers work
+    // cross-platform.
     for (name, template) in ACTIVATE_TEMPLATES {
         let activator = template
             .replace("{{ VIRTUAL_ENV_DIR }}", location.as_str())

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

@@ -67,8 +67,8 @@ pub enum Error {
     RecordCsv(#[from] csv::Error),
     #[error("Broken virtualenv: {0}")]
     BrokenVenv(String),
-    #[error("Failed to detect the operating system version: {0}")]
-    OsVersionDetection(String),
+    #[error("Unable to create Windows launch for {0} (only x64_64 is supported)")]
+    UnsupportedWindowsArch(&'static str),
     #[error("Failed to detect the current platform")]
     PlatformInfo(#[source] PlatformInfoError),
     #[error("Invalid version specification, only none or == is supported")]

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

@@ -99,8 +99,14 @@ pub fn install_wheel(
 
     debug!(name, "Writing entrypoints");
     let (console_scripts, gui_scripts) = parse_scripts(&wheel, &dist_info_prefix, None)?;
-    write_script_entrypoints(&site_packages, location, &console_scripts, &mut record)?;
-    write_script_entrypoints(&site_packages, location, &gui_scripts, &mut record)?;
+    write_script_entrypoints(
+        &site_packages,
+        location,
+        &console_scripts,
+        &mut record,
+        false,
+    )?;
+    write_script_entrypoints(&site_packages, location, &gui_scripts, &mut record, true)?;
 
     let data_dir = site_packages.join(format!("{dist_info_prefix}.data"));
     // 2.a Unpacked archive includes distribution-1.0.dist-info/ and (if there is data) distribution-1.0.data/.

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

@@ -34,9 +34,10 @@ use crate::{find_dist_info, Error};
 /// `#!/usr/bin/env python`
 pub const SHEBANG_PYTHON: &str = "#!/usr/bin/env python";
 
-pub(crate) const LAUNCHER_T32: &[u8] = include_bytes!("../windows-launcher/t32.exe");
-pub(crate) const LAUNCHER_T64: &[u8] = include_bytes!("../windows-launcher/t64.exe");
-pub(crate) const LAUNCHER_T64_ARM: &[u8] = include_bytes!("../windows-launcher/t64-arm.exe");
+pub(crate) const LAUNCHER_X86_64_GUI: &[u8] =
+    include_bytes!("../../puffin-trampoline/trampolines/puffin-trampoline-gui.exe");
+pub(crate) const LAUNCHER_X86_64_CONSOLE: &[u8] =
+    include_bytes!("../../puffin-trampoline/trampolines/puffin-trampoline-console.exe");
 
 /// Wrapper script template function
 ///
@@ -283,35 +284,34 @@ pub(crate) fn get_shebang(location: &InstallLocation<impl AsRef<Path>>) -> Strin
     format!("#!{path}")
 }
 
-/// To get a launcher on windows we write a minimal .exe launcher binary and then attach the actual
-/// python after it.
-///
-/// TODO pyw scripts
-///
-/// TODO: a nice, reproducible-without-distlib rust solution
+/// A windows script is a minimal .exe launcher binary with the python entrypoint script appended as stored zip file.
+/// The launcher will look for `python[w].exe` adjacent to it in the same directory to start the embedded script.
 ///
 /// <https://github.com/pypa/pip/blob/fd0ea6bc5e8cb95e518c23d901c26ca14db17f89/src/pip/_vendor/distlib/scripts.py#L248-L262>
-pub(crate) fn windows_script_launcher(launcher_python_script: &str) -> Result<Vec<u8>, Error> {
+pub(crate) fn windows_script_launcher(
+    launcher_python_script: &str,
+    is_gui: bool,
+) -> Result<Vec<u8>, Error> {
     let launcher_bin = match env::consts::ARCH {
-        "x84" => LAUNCHER_T32,
-        "x86_64" => LAUNCHER_T64,
-        "aarch64" => LAUNCHER_T64_ARM,
+        "x86_64" => {
+            if is_gui {
+                LAUNCHER_X86_64_GUI
+            } else {
+                LAUNCHER_X86_64_CONSOLE
+            }
+        }
         arch => {
-            let error = format!(
-                "Don't know how to create windows launchers for script for {arch}, \
-                        only x86, x86_64 and aarch64 (64-bit arm) are supported"
-            );
-            return Err(Error::OsVersionDetection(error));
+            return Err(Error::UnsupportedWindowsArch(arch));
         }
     };
 
-    let mut stream: Vec<u8> = Vec::new();
+    let mut payload: Vec<u8> = Vec::new();
     {
-        // We're using the zip writer, but it turns out we're not actually deflating apparently
-        // we're just using an offset
+        // We're using the zip writer, but with stored compression
+        // https://github.com/njsmith/posy/blob/04927e657ca97a5e35bb2252d168125de9a3a025/src/trampolines/mod.rs#L75-L82
         // https://github.com/pypa/distlib/blob/8ed03aab48add854f377ce392efffb79bb4d6091/PC/launcher.c#L259-L271
         let stored = FileOptions::default().compression_method(zip::CompressionMethod::Stored);
-        let mut archive = ZipWriter::new(Cursor::new(&mut stream));
+        let mut archive = ZipWriter::new(Cursor::new(&mut payload));
         let error_msg = "Writing to Vec<u8> should never fail";
         archive.start_file("__main__.py", stored).expect(error_msg);
         archive
@@ -320,8 +320,9 @@ pub(crate) fn windows_script_launcher(launcher_python_script: &str) -> Result<Ve
         archive.finish().expect(error_msg);
     }
 
-    let mut launcher: Vec<u8> = launcher_bin.to_vec();
-    launcher.append(&mut stream);
+    let mut launcher: Vec<u8> = Vec::with_capacity(launcher_bin.len() + payload.len());
+    launcher.extend_from_slice(launcher_bin);
+    launcher.extend_from_slice(&payload);
     Ok(launcher)
 }
 
@@ -335,6 +336,7 @@ pub(crate) fn write_script_entrypoints(
     location: &InstallLocation<impl AsRef<Path>>,
     entrypoints: &[Script],
     record: &mut Vec<RecordEntry>,
+    is_gui: bool,
 ) -> Result<(), Error> {
     for entrypoint in entrypoints {
         let entrypoint_relative = if cfg!(windows) {
@@ -356,7 +358,7 @@ pub(crate) fn write_script_entrypoints(
             &get_shebang(location),
         );
         if cfg!(windows) {
-            let launcher = windows_script_launcher(&launcher_python_script)?;
+            let launcher = windows_script_launcher(&launcher_python_script, is_gui)?;
             write_file_recorded(site_packages, &entrypoint_relative, &launcher, record)?;
         } else {
             write_file_recorded(
@@ -1009,8 +1011,14 @@ pub fn install_wheel(
 
     debug!(name = name.as_str(), "Writing entrypoints");
     let (console_scripts, gui_scripts) = parse_scripts(&mut archive, &dist_info_prefix, None)?;
-    write_script_entrypoints(&site_packages, location, &console_scripts, &mut record)?;
-    write_script_entrypoints(&site_packages, location, &gui_scripts, &mut record)?;
+    write_script_entrypoints(
+        &site_packages,
+        location,
+        &console_scripts,
+        &mut record,
+        false,
+    )?;
+    write_script_entrypoints(&site_packages, location, &gui_scripts, &mut record, true)?;
 
     let data_dir = site_packages.join(format!("{dist_info_prefix}.data"));
     // 2.a Unpacked archive includes distribution-1.0.dist-info/ and (if there is data) distribution-1.0.data/.
@@ -1135,7 +1143,9 @@ mod test {
 
     use indoc::{formatdoc, indoc};
 
-    use crate::wheel::{read_record_file, relative_to};
+    use crate::wheel::{
+        read_record_file, relative_to, LAUNCHER_X86_64_CONSOLE, LAUNCHER_X86_64_GUI,
+    };
     use crate::{parse_key_value_file, Script};
 
     use super::parse_wheel_version;
@@ -1264,4 +1274,19 @@ mod test {
             })
         );
     }
+
+    #[test]
+    fn test_launchers_are_small() {
+        // At time of writing, they are 15872 bytes.
+        assert!(
+            LAUNCHER_X86_64_GUI.len() < 20 * 1024,
+            "GUI launcher: {}",
+            LAUNCHER_X86_64_GUI.len()
+        );
+        assert!(
+            LAUNCHER_X86_64_CONSOLE.len() < 20 * 1024,
+            "CLI launcher: {}",
+            LAUNCHER_X86_64_CONSOLE.len()
+        );
+    }
 }