Skip to content

Commit

Permalink
Allow relative Python executable paths in Windows trampoline (#3717)
Browse files Browse the repository at this point in the history
## Summary

This is a prerequisite for #3669

## Test Plan

Download one of the standalone distributions on Windows then use its
Python to run the following script and then run the scripts it creates
(only pip):

```python
from __future__ import annotations

import sys
import sysconfig
from contextlib import closing
from importlib.metadata import entry_points
from io import BytesIO
from os.path import relpath
from pathlib import Path
from tempfile import TemporaryDirectory
from zipfile import ZIP_DEFLATED, ZipFile, ZipInfo

# Change this line to your real path
LAUNCHERS_DIR = Path('C:\\Users\\ofek\\Desktop\\code\\uv\\crates\\uv-trampoline\\target\\x86_64-pc-windows-msvc\\release')
SCRIPT_TEMPLATE = """\
#!{executable}
# -*- coding: utf-8 -*-
import re
import sys
from {module} import {import_name}
if __name__ == "__main__":
    sys.argv[0] = re.sub(r"(-script\\.pyw|\\.exe)?$", "", sys.argv[0])
    sys.exit({function}())
"""


def select_entry_points(ep, group):
    return ep.select(group=group) if sys.version_info[:2] >= (3, 10) else ep.get(group, [])


def main():
    interpreters_dir = Path(sys.executable).parent
    scripts_dir = Path(sysconfig.get_path('scripts'))

    ep = entry_points()
    for group, interpreter_name, launcher_name in (
        ('console_scripts', 'python.exe', 'uv-trampoline-console.exe'),
        ('gui_scripts', 'pythonw.exe', 'uv-trampoline-gui.exe'),
    ):
        interpreter = interpreters_dir / interpreter_name
        relative_interpreter_path = relpath(interpreter, scripts_dir)
        launcher_data = (LAUNCHERS_DIR / launcher_name).read_bytes()

        for script in select_entry_points(ep, group):
            # https://github.com/astral-sh/uv/tree/main/crates/uv-trampoline#how-do-you-use-it
            with closing(BytesIO()) as buf:
                # Launcher
                buf.write(launcher_data)

                # Zipped script
                with TemporaryDirectory() as td:
                    zip_path = Path(td) / 'script.zip'
                    with ZipFile(zip_path, 'w') as zf:
                        # Ensure reproducibility
                        zip_info = ZipInfo('__main__.py', (2020, 2, 2, 0, 0, 0))
                        zip_info.external_attr = (0o644 & 0xFFFF) << 16

                        module, _, attrs = script.value.partition(':')
                        contents = SCRIPT_TEMPLATE.format(
                            executable=relative_interpreter_path,
                            module=module,
                            import_name=attrs.split('.')[0],
                            function=attrs
                        )
                        zf.writestr(zip_info, contents, compress_type=ZIP_DEFLATED)

                    buf.write(zip_path.read_bytes())

                # Interpreter path
                interpreter_path = relative_interpreter_path.encode('utf-8')
                buf.write(interpreter_path)

                # Interpreter path length
                interpreter_path_length = len(interpreter_path).to_bytes(4, 'little')
                buf.write(interpreter_path_length)

                # Magic number
                buf.write(b'UVUV')

                script_data = buf.getvalue()

            script_path = scripts_dir / f'{script.name}.exe'
            script_path.write_bytes(script_data)


if __name__ == '__main__':
    main()
```
  • Loading branch information
ofek authored May 22, 2024
1 parent becdc64 commit 82820d0
Show file tree
Hide file tree
Showing 5 changed files with 31 additions and 3 deletions.
34 changes: 31 additions & 3 deletions crates/uv-trampoline/src/bounce.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ fn make_child_cmdline() -> CString {

child_cmdline.push(b'\0');

// Helpful when debugging trampline issues
// Helpful when debugging trampoline issues
// eprintln!(
// "executable_name: '{}'\nnew_cmdline: {}",
// core::str::from_utf8(executable_name.to_bytes()).unwrap(),
Expand Down Expand Up @@ -183,7 +183,7 @@ fn find_python_exe(executable_name: &CStr) -> CString {
let mut buffer: Vec<u8> = Vec::new();
let mut bytes_to_read = 1024.min(u32::try_from(file_size).unwrap_or(u32::MAX));

let path = loop {
let path: CString = loop {
// SAFETY: Casting to usize is safe because we only support 64bit systems where usize is guaranteed to be larger than u32.
buffer.resize(bytes_to_read as usize, 0);

Expand Down Expand Up @@ -275,7 +275,35 @@ fn find_python_exe(executable_name: &CStr) -> CString {
String::from("Failed to close file handle")
});

path
if is_absolute(&path) {
path
} else {
let parent_dir = match executable_name
.to_bytes()
.rsplitn(2, |c| *c == b'\\')
.last()
{
Some(parent_dir) => parent_dir,
None => {
eprintln!("Script path has unknown separator characters.");
exit_with_status(1)
}
};
let final_path = [parent_dir, b"\\", path.as_bytes()].concat();
CString::new(final_path).unwrap_or_else(|_| {
eprintln!("Could not construct the absolute path to the Python executable.");
exit_with_status(1)
})
}
}

/// Returns `true` if the path is absolute.
///
/// In this context, as in the Rust standard library, c:\windows is absolute, while c:temp and
/// \temp are not.
fn is_absolute(path: &CStr) -> bool {
let path = path.to_bytes();
path.len() >= 3 && path[0].is_ascii_alphabetic() && path[1] == b':' && path[2] == b'\\'
}

fn push_arguments(output: &mut Vec<u8>) {
Expand Down
Binary file modified crates/uv-trampoline/trampolines/uv-trampoline-aarch64-console.exe
Binary file not shown.
Binary file modified crates/uv-trampoline/trampolines/uv-trampoline-aarch64-gui.exe
Binary file not shown.
Binary file modified crates/uv-trampoline/trampolines/uv-trampoline-x86_64-console.exe
Binary file not shown.
Binary file not shown.

0 comments on commit 82820d0

Please sign in to comment.