From 57338e558cdf7f0e82a49ae5530439016e43a450 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 9 Jul 2025 11:51:06 -0400 Subject: [PATCH] Drop trailing arguments when writing shebangs (#14519) ## Summary You can see in pip that they read the full first line, then replace it with the rewritten shebang, thereby dropping any trailing arguments on the shebang: https://github.com/pypa/pip/blob/65da0ff5349297da64ccadb4dd22ab41185ea0b9/src/pip/_internal/operations/install/wheel.py#L94 In contrast, we currently retain them, but write them _after_ the shebang, which is wrong. Closes https://github.com/astral-sh/uv/issues/14470. --- crates/uv-install-wheel/src/wheel.rs | 43 ++++++++--- crates/uv/tests/it/pip_install.rs | 107 +++++++++++++++++++++++++++ 2 files changed, 138 insertions(+), 12 deletions(-) diff --git a/crates/uv-install-wheel/src/wheel.rs b/crates/uv-install-wheel/src/wheel.rs index d013ac5eb..250143016 100644 --- a/crates/uv-install-wheel/src/wheel.rs +++ b/crates/uv-install-wheel/src/wheel.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; use std::io; -use std::io::{BufReader, Read, Seek, Write}; +use std::io::{BufReader, Read, Write}; use std::path::{Path, PathBuf}; use data_encoding::BASE64URL_NOPAD; @@ -144,7 +144,7 @@ fn format_shebang(executable: impl AsRef, os_name: &str, relocatable: bool /// /// fn get_script_executable(python_executable: &Path, is_gui: bool) -> PathBuf { - // Only check for pythonw.exe on Windows + // Only check for `pythonw.exe` on Windows. if cfg!(windows) && is_gui { python_executable .file_name() @@ -431,22 +431,41 @@ fn install_script( Err(err) => return Err(Error::Io(err)), } let size_and_encoded_hash = if start == placeholder_python { - let is_gui = { - let mut buf = vec![0; 1]; - script.read_exact(&mut buf)?; - if buf == b"w" { - true - } else { - script.seek_relative(-1)?; - false + // Read the rest of the first line, one byte at a time, until we hit a newline. + let mut is_gui = false; + let mut first = true; + let mut byte = [0u8; 1]; + loop { + match script.read_exact(&mut byte) { + Ok(()) => { + if byte[0] == b'\n' || byte[0] == b'\r' { + break; + } + + // Check if this is a GUI script (starts with 'w'). + if first { + is_gui = byte[0] == b'w'; + first = false; + } + } + Err(err) if err.kind() == io::ErrorKind::UnexpectedEof => break, + Err(err) => return Err(Error::Io(err)), } - }; + } + let executable = get_script_executable(&layout.sys_executable, is_gui); let executable = get_relocatable_executable(executable, layout, relocatable)?; - let start = format_shebang(&executable, &layout.os_name, relocatable) + let mut start = format_shebang(&executable, &layout.os_name, relocatable) .as_bytes() .to_vec(); + // Use appropriate line ending for the platform. + if layout.os_name == "nt" { + start.extend_from_slice(b"\r\n"); + } else { + start.push(b'\n'); + } + let mut target = uv_fs::tempfile_in(&layout.scheme.scripts)?; let size_and_encoded_hash = copy_and_hash(&mut start.chain(script), &mut target)?; diff --git a/crates/uv/tests/it/pip_install.rs b/crates/uv/tests/it/pip_install.rs index a33e08d90..f231198e4 100644 --- a/crates/uv/tests/it/pip_install.rs +++ b/crates/uv/tests/it/pip_install.rs @@ -11490,3 +11490,110 @@ fn conflicting_flags_clap_bug() { " ); } + +/// Test that shebang arguments are stripped when installing scripts +#[test] +#[cfg(unix)] +fn strip_shebang_arguments() -> Result<()> { + let context = TestContext::new("3.12"); + + let project_dir = context.temp_dir.child("shebang_test"); + project_dir.create_dir_all()?; + + // Create a package with scripts that have shebang arguments. + let pyproject_toml = project_dir.child("pyproject.toml"); + pyproject_toml.write_str(indoc! {r#" + [project] + name = "shebang-test" + version = "0.1.0" + + [build-system] + requires = ["setuptools>=61.0"] + build-backend = "setuptools.build_meta" + + [tool.setuptools] + packages = ["shebang_test"] + + [tool.setuptools.data-files] + "scripts" = ["scripts/custom_script", "scripts/custom_gui_script"] + "#})?; + + // Create the package directory. + let package_dir = project_dir.child("shebang_test"); + package_dir.create_dir_all()?; + + // Create an `__init__.py` file in the package directory. + let init_file = package_dir.child("__init__.py"); + init_file.touch()?; + + // Create scripts directory with scripts that have shebangs with arguments + let scripts_dir = project_dir.child("scripts"); + scripts_dir.create_dir_all()?; + + let script_with_args = scripts_dir.child("custom_script"); + script_with_args.write_str(indoc! {r#" + #!python -E -s + # This is a test script with shebang arguments + import sys + print(f"Hello from {sys.executable}") + print(f"Arguments: {sys.argv}") + "#})?; + + let gui_script_with_args = scripts_dir.child("custom_gui_script"); + gui_script_with_args.write_str(indoc! {r#" + #!pythonw -E + # This is a test GUI script with shebang arguments + import sys + print(f"Hello from GUI script: {sys.executable}") + "#})?; + + // Create a `setup.py` that explicitly handles scripts. + let setup_py = project_dir.child("setup.py"); + setup_py.write_str(indoc! {r" + from setuptools import setup + setup(scripts=['scripts/custom_script', 'scripts/custom_gui_script']) + "})?; + + // Install the package. + uv_snapshot!(context.filters(), context.pip_install().arg(project_dir.path()), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Resolved 1 package in [TIME] + Prepared 1 package in [TIME] + Installed 1 package in [TIME] + + shebang-test==0.1.0 (from file://[TEMP_DIR]/shebang_test) + "###); + + // Check the installed scripts have their shebangs stripped of arguments. + let custom_script_path = venv_bin_path(&context.venv).join("custom_script"); + let script_content = fs::read_to_string(&custom_script_path)?; + + insta::with_settings!({filters => context.filters() + }, { + insta::assert_snapshot!(script_content, @r#" + #![VENV]/bin/python3 + # This is a test script with shebang arguments + import sys + print(f"Hello from {sys.executable}") + print(f"Arguments: {sys.argv}") + "#); + }); + + let custom_gui_script_path = venv_bin_path(&context.venv).join("custom_gui_script"); + let gui_script_content = fs::read_to_string(&custom_gui_script_path)?; + + insta::with_settings!({filters => context.filters() + }, { + insta::assert_snapshot!(gui_script_content, @r#" + #![VENV]/bin/python3 + # This is a test GUI script with shebang arguments + import sys + print(f"Hello from GUI script: {sys.executable}") + "#); + }); + + Ok(()) +}