Lookup Python executable from binary

Author: MichaReiser

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

@@ -33,6 +33,8 @@ 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");
@@ -235,7 +237,7 @@ fn unpack_wheel_files<R: Read + Seek>(
             record_path.with_extension("jws"),
             record_path.with_extension("p7s"),
         ]
-        .contains(&relative)
+            .contains(&relative)
         {
             continue;
         }
@@ -281,18 +283,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}")
 }
 
@@ -305,6 +296,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.
@@ -352,9 +344,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)
 }
 
@@ -393,7 +396,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 {
@@ -949,7 +952,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,
@@ -1224,8 +1227,8 @@ mod test {
             "selenium-4.1.0.dist-info/METADATA",
             "selenium-4.1.0.dist-info/RECORD",
         ]
-        .map(ToString::to_string)
-        .to_vec();
+            .map(ToString::to_string)
+            .to_vec();
         let actual = entries
             .into_iter()
             .map(|entry| entry.path)
@@ -1240,23 +1243,23 @@ mod test {
                 Path::new("/home/ferris/carcinization/lib/python/site-packages/foo/__init__.py"),
                 Path::new("/home/ferris/carcinization/lib/python/site-packages"),
             )
-            .unwrap(),
+                .unwrap(),
             Path::new("foo/__init__.py")
         );
         assert_eq!(
             relative_to(
                 Path::new("/home/ferris/carcinization/lib/marker.txt"),
                 Path::new("/home/ferris/carcinization/lib/python/site-packages"),
             )
-            .unwrap(),
+                .unwrap(),
             Path::new("../../marker.txt")
         );
         assert_eq!(
             relative_to(
                 Path::new("/home/ferris/carcinization/bin/foo_launcher"),
                 Path::new("/home/ferris/carcinization/lib/python/site-packages"),
             )
-            .unwrap(),
+                .unwrap(),
             Path::new("../../../bin/foo_launcher")
         );
     }
@@ -1277,7 +1280,7 @@ mod test {
                 "foo.bar:main",
                 Some(&["bar".to_string(), "baz".to_string()]),
             )
-            .unwrap(),
+                .unwrap(),
             Some(Script {
                 script_name: "launcher".to_string(),
                 module: "foo.bar".to_string(),
@@ -1294,7 +1297,7 @@ mod test {
                 "foomod:main_bar [bar,baz]",
                 Some(&["bar".to_string(), "baz".to_string()]),
             )
-            .unwrap(),
+                .unwrap(),
             Some(Script {
                 script_name: "launcher".to_string(),
                 module: "foomod".to_string(),

crates/uv-trampoline/Cargo.lock

@@ -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?
 
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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.