Added ability to select bytecode invalidation mode of generated .pyc

crates/uv-installer/src/compile.rs

@@ -32,7 +32,7 @@ pub enum CompileError {
     PythonSubcommand(#[source] io::Error),
     #[error("Failed to create temporary script file")]
     TempFile(#[source] io::Error),
-    #[error("Bytecode compilation failed, expected {0:?}, received: {1:?}")]
+    #[error(r#"Bytecode compilation failed, expected "{0}", received: "{1}""#)]
     WrongPath(String, String),
     #[error("Failed to write to Python {device}")]
     ChildStdio {
@@ -82,7 +82,7 @@ pub async fn compile_tree(
     let tempdir = tempdir_in(cache).map_err(CompileError::TempFile)?;
     let pip_compileall_py = tempdir.path().join("pip_compileall.py");
 
-    // Start the workers.
+    debug!("Starting {} bytecode compilation workers", worker_count);
     let mut worker_handles = Vec::new();
     for _ in 0..worker_count.get() {
         worker_handles.push(tokio::task::spawn(worker(
@@ -92,6 +92,8 @@ pub async fn compile_tree(
             receiver.clone(),
         )));
     }
+    // Make sure the channel gets closed when all workers exit.
+    drop(receiver);
 
     // Start the producer, sending all `.py` files to workers.
     let mut source_files = 0;
@@ -191,9 +193,11 @@ async fn worker(
             device: "stderr",
             err,
         })?;
-    if !child_stderr_collected.is_empty() {
+    let result = if child_stderr_collected.is_empty() {
+        result
+    } else {
         let stderr = String::from_utf8_lossy(&child_stderr_collected);
-        return match result {
+        match result {
             Ok(()) => {
                 debug!(
                     "Bytecode compilation `python` at {} stderr:\n{}\n---",
@@ -203,11 +207,13 @@ async fn worker(
                 Ok(())
             }
             Err(err) => Err(CompileError::ErrorWithStderr {
-                stderr: stderr.to_string(),
+                stderr: stderr.trim().to_string(),
                 err: Box::new(err),
             }),
-        };
-    }
+        }
+    };
+
+    debug!("Bytecode compilation worker exiting: {:?}", result);
 
     result
 }

crates/uv-installer/src/pip_compileall.py

@@ -9,6 +9,8 @@
 """
 
 import compileall
+import os
+import py_compile
 import sys
 import warnings
 
@@ -18,6 +20,23 @@
     # Successful launch check
     print("Ready")
 
+    # https://docs.python.org/3/library/py_compile.html#py_compile.PycInvalidationMode
+    # TIMESTAMP, CHECKED_HASH, UNCHECKED_HASH
+    invalidation_mode = os.environ.get("PYC_INVALIDATION_MODE")
+    if invalidation_mode is not None:
+        try:
+            invalidation_mode = py_compile.PycInvalidationMode[invalidation_mode]
+        except KeyError:
+            invalidation_modes = ", ".join(
+                '"' + x.name + '"' for x in py_compile.PycInvalidationMode
+            )
+            print(
+                f'Invalid value for PYC_INVALIDATION_MODE "{invalidation_mode}", '
+                f"valid are {invalidation_modes}: ",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+
     # In rust, we provide one line per file to compile.
     for path in sys.stdin:
         # Remove trailing newlines.
@@ -27,6 +46,8 @@
         # Unlike pip, we set quiet=2, so we don't have to capture stdout.
         # We'd like to show those errors, but given that pip thinks that's totally fine,
         # we can't really change that.
-        success = compileall.compile_file(path, force=True, quiet=2)
+        success = compileall.compile_file(
+            path, invalidation_mode=invalidation_mode, force=True, quiet=2
+        )
         # We're ready for the next file.
         print(path)

crates/uv/tests/common/mod.rs

@@ -174,6 +174,20 @@ impl TestContext {
     pub fn python_kind(&self) -> &str {
         "python"
     }
+
+    /// Returns the site-packages folder inside the venv.
+    pub fn site_packages(&self) -> PathBuf {
+        if cfg!(unix) {
+            self.venv
+                .join("lib")
+                .join(format!("{}{}", self.python_kind(), self.python_version))
+                .join("site-packages")
+        } else if cfg!(windows) {
+            self.venv.join("Lib").join("site-packages")
+        } else {
+            unimplemented!("Only Windows and Unix are supported")
+        }
+    }
 }
 
 pub fn venv_to_interpreter(venv: &Path) -> PathBuf {

crates/uv/tests/pip_sync.rs

@@ -2,7 +2,7 @@
 
 use fs_err as fs;
 use std::env::consts::EXE_SUFFIX;
-use std::path::{Path, PathBuf};
+use std::path::Path;
 use std::process::Command;
 
 use anyhow::Result;
@@ -73,25 +73,6 @@ fn uninstall_command(context: &TestContext) -> Command {
     command
 }
 
-/// Returns the site-packages folder inside the venv.
-fn site_packages(context: &TestContext) -> PathBuf {
-    if cfg!(unix) {
-        context
-            .venv
-            .join("lib")
-            .join(format!(
-                "{}{}",
-                context.python_kind(),
-                context.python_version
-            ))
-            .join("site-packages")
-    } else if cfg!(windows) {
-        context.venv.join("Lib").join("site-packages")
-    } else {
-        unimplemented!("Only Windows and Unix are supported")
-    }
-}
-
 #[test]
 fn missing_pip() {
     uv_snapshot!(Command::new(get_bin()).arg("sync"), @r###"
@@ -186,7 +167,8 @@ fn install() -> Result<()> {
     );
 
     // Counterpart for the `compile()` test.
-    assert!(!site_packages(&context)
+    assert!(!context
+        .site_packages()
         .join("markupsafe")
         .join("__pycache__")
         .join("__init__.cpython-312.pyc")
@@ -2946,7 +2928,8 @@ fn compile() -> Result<()> {
     "###
     );
 
-    assert!(site_packages(&context)
+    assert!(context
+        .site_packages()
         .join("markupsafe")
         .join("__pycache__")
         .join("__init__.cpython-312.pyc")
@@ -2957,6 +2940,57 @@ fn compile() -> Result<()> {
     Ok(())
 }
 
+/// Test that the `PYC_INVALIDATION_MODE` option is recognized and that the error handling works.
+#[test]
+fn compile_invalid_pyc_invalidation_mode() -> Result<()> {
+    let context = TestContext::new("3.12");
+
+    let requirements_txt = context.temp_dir.child("requirements.txt");
+    requirements_txt.touch()?;
+    requirements_txt.write_str("MarkupSafe==2.1.3")?;
+
+    let site_packages = regex::escape(
+        &context
+            .site_packages()
+            .canonicalize()
+            .unwrap()
+            .simplified_display()
+            .to_string(),
+    );
+    let filters: Vec<_> = [
+        (site_packages.as_str(), "[SITE-PACKAGES]"),
+        (
+            r#"\[SITE-PACKAGES\].*.py", received: "#,
+            r#"[SITE-PACKAGES]/[FIRST-FILE]", received: "#,
+        ),
+    ]
+    .into_iter()
+    .chain(INSTA_FILTERS.to_vec())
+    .collect();
+
+    uv_snapshot!(filters, command(&context)
+        .arg("requirements.txt")
+        .arg("--compile")
+        .arg("--strict")
+        .env("PYC_INVALIDATION_MODE", "bogus"), @r###"
+    success: false
+    exit_code: 2
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 1 package in [TIME]
+    Downloaded 1 package in [TIME]
+    Installed 1 package in [TIME]
+    error: Failed to bytecode compile [SITE-PACKAGES]
+      Caused by: Python process stderr:
+    Invalid value for PYC_INVALIDATION_MODE "bogus", valid are "TIMESTAMP", "CHECKED_HASH", "UNCHECKED_HASH":
+      Caused by: Bytecode compilation failed, expected "[SITE-PACKAGES]/[FIRST-FILE]", received: ""
+    "###
+    );
+
+    Ok(())
+}
+
 /// Raise an error when an editable's `Requires-Python` constraint is not met.
 #[test]
 fn requires_python_editable() -> Result<()> {