Merge 9693f2ef91 into 0a83bf7dd5

crates/uv-cli/src/lib.rs

@@ -6112,6 +6112,37 @@ pub struct PythonDirArgs {
     pub bin: bool,
 }
 
+#[derive(Args)]
+pub struct PythonInstallCompileBytecodeArgs {
+    /// Compile Python's standard library to bytecode after installation.
+    ///
+    /// By default, uv does not compile Python (`.py`) files to bytecode (`__pycache__/*.pyc`);
+    /// instead, compilation is performed lazily the first time a module is imported. For use-cases
+    /// in which start time is critical, such as CLI applications and Docker containers, this option
+    /// can be enabled to trade longer installation times for faster start times.
+    ///
+    /// When enabled, uv will process the Python version's `stdlib` directory. Like pip, it will
+    /// also ignore errors.
+    #[arg(
+        long,
+        alias = "compile",
+        overrides_with("no_compile_bytecode"),
+        help_heading = "Installer options",
+        env = EnvVars::UV_COMPILE_BYTECODE,
+        value_parser = clap::builder::BoolishValueParser::new(),
+    )]
+    pub compile_bytecode: bool,
+
+    #[arg(
+        long,
+        alias = "no-compile",
+        overrides_with("compile_bytecode"),
+        hide = true,
+        help_heading = "Installer options"
+    )]
+    pub no_compile_bytecode: bool,
+}
+
 #[derive(Args)]
 pub struct PythonInstallArgs {
     /// The directory to store the Python installation in.
@@ -6231,6 +6262,9 @@ pub struct PythonInstallArgs {
     /// If multiple Python versions are requested, uv will exit with an error.
     #[arg(long, conflicts_with("no_bin"))]
     pub default: bool,
+
+    #[command(flatten)]
+    pub compile_bytecode: PythonInstallCompileBytecodeArgs,
 }
 
 impl PythonInstallArgs {
@@ -6291,6 +6325,9 @@ pub struct PythonUpgradeArgs {
     /// URL pointing to JSON of custom Python installations.
     #[arg(long, value_hint = ValueHint::Other)]
     pub python_downloads_json_url: Option<String>,
+
+    #[command(flatten)]
+    pub compile_bytecode: PythonInstallCompileBytecodeArgs,
 }
 
 impl PythonUpgradeArgs {

crates/uv/src/commands/python/install.rs

@@ -5,16 +5,18 @@ use std::io::ErrorKind;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
 
-use anyhow::{Error, Result};
+use anyhow::{Context, Error, Result};
-use futures::StreamExt;
+use futures::{StreamExt, join};
-use futures::stream::FuturesUnordered;
 use indexmap::IndexSet;
 use itertools::{Either, Itertools};
 use owo_colors::{AnsiColors, OwoColorize};
 use rustc_hash::{FxHashMap, FxHashSet};
+use tokio::sync::mpsc;
 use tracing::{debug, trace};
 
+use uv_cache::Cache;
 use uv_client::BaseClientBuilder;
+use uv_configuration::Concurrency;
 use uv_fs::Simplified;
 use uv_platform::{Arch, Libc};
 use uv_preview::{Preview, PreviewFeatures};
@@ -27,12 +29,13 @@ use uv_python::managed::{
     create_link_to_executable, python_executable_dir,
 };
 use uv_python::{
-    PythonDownloads, PythonInstallationKey, PythonInstallationMinorVersionKey, PythonRequest,
+    Interpreter, PythonDownloads, PythonInstallationKey, PythonInstallationMinorVersionKey,
-    PythonVersionFile, VersionFileDiscoveryOptions, VersionFilePreference, VersionRequest,
+    PythonRequest, PythonVersionFile, VersionFileDiscoveryOptions, VersionFilePreference,
+    VersionRequest,
 };
 use uv_shell::Shell;
 use uv_trampoline_builder::{Launcher, LauncherKind};
-use uv_warnings::{warn_user, write_error_chain};
+use uv_warnings::{warn_user, warn_user_once, write_error_chain};
 
 use crate::commands::python::{ChangeEvent, ChangeEventKind};
 use crate::commands::reporters::PythonDownloadReporter;
@@ -191,6 +194,93 @@ pub(crate) async fn install(
     default: bool,
     python_downloads: PythonDownloads,
     no_config: bool,
+    compile_bytecode: bool,
+    concurrency: &Concurrency,
+    cache: &Cache,
+    preview: Preview,
+    printer: Printer,
+) -> Result<ExitStatus> {
+    let (sender, mut receiver) = mpsc::unbounded_channel();
+    let compiler = async {
+        let mut did_compile = false;
+        let mut total_files = 0;
+        let mut total_elapsed = std::time::Duration::default();
+        while let Some(installation) = receiver.recv().await {
+            did_compile = true;
+            let (files, elapsed) = compile_stdlib_bytecode(&installation, concurrency, cache)
+                .await
+                .with_context(|| {
+                    format!(
+                        "Failed to bytecode-compile Python standard library for: {}",
+                        installation.key()
+                    )
+                })?;
+            total_files += files;
+            total_elapsed += elapsed;
+        }
+        Ok::<_, anyhow::Error>(did_compile.then_some((total_files, total_elapsed)))
+    };
+
+    let installer = perform_install(
+        project_dir,
+        install_dir,
+        targets,
+        reinstall,
+        upgrade,
+        bin,
+        registry,
+        force,
+        python_install_mirror,
+        pypy_install_mirror,
+        python_downloads_json_url,
+        client_builder,
+        default,
+        python_downloads,
+        no_config,
+        compile_bytecode.then_some(sender),
+        concurrency,
+        preview,
+        printer,
+    );
+
+    let (installer_result, compiler_result) = join!(installer, compiler);
+
+    if let Some((total_files, total_elapsed)) = compiler_result? {
+        let s = if total_files == 1 { "" } else { "s" };
+        writeln!(
+            printer.stderr(),
+            "{}",
+            format!(
+                "Bytecode compiled {} {}",
+                format!("{total_files} file{s}").bold(),
+                format!("in {}", elapsed(total_elapsed)).dimmed(),
+            )
+            .dimmed()
+        )?;
+    }
+
+    installer_result
+}
+
+#[allow(clippy::fn_params_excessive_bools)]
+async fn perform_install(
+    project_dir: &Path,
+    install_dir: Option<PathBuf>,
+    targets: Vec<String>,
+    reinstall: bool,
+    upgrade: PythonUpgrade,
+    bin: Option<bool>,
+    registry: Option<bool>,
+    force: bool,
+    python_install_mirror: Option<String>,
+    pypy_install_mirror: Option<String>,
+    python_downloads_json_url: Option<String>,
+    client_builder: BaseClientBuilder<'_>,
+    default: bool,
+    python_downloads: PythonDownloads,
+    no_config: bool,
+    bytecode_compilation_sender: Option<mpsc::UnboundedSender<ManagedPythonInstallation>>,
+    concurrency: &Concurrency,
     preview: Preview,
     printer: Printer,
 ) -> Result<ExitStatus> {
@@ -433,6 +523,16 @@ pub(crate) async fn install(
         })
     };
 
+    // For all satisfied installs, bytecode compile them now before any future
+    // early return.
+    if let Some(ref sender) = bytecode_compilation_sender {
+        satisfied
+            .iter()
+            .copied()
+            .cloned()
+            .try_for_each(|installation| sender.send(installation))?;
+    }
+
     // Check if Python downloads are banned
     if matches!(python_downloads, PythonDownloads::Never) && !unsatisfied.is_empty() {
         writeln!(
@@ -458,10 +558,9 @@ pub(crate) async fn install(
 
     // Download and unpack the Python versions concurrently
     let reporter = PythonDownloadReporter::new(printer, Some(downloads.len() as u64));
-    let mut tasks = FuturesUnordered::new();
 
-    for download in &downloads {
+    let mut tasks = futures::stream::iter(&downloads)
-        tasks.push(async {
+        .map(async |download| {
             (
                 *download,
                 download
@@ -477,8 +576,8 @@ pub(crate) async fn install(
                     )
                     .await,
             )
-        });
+        })
-    }
+        .buffer_unordered(concurrency.downloads);
 
     let mut errors = vec![];
     let mut downloaded = Vec::with_capacity(downloads.len());
@@ -493,6 +592,9 @@ pub(crate) async fn install(
                 };
 
                 let installation = ManagedPythonInstallation::new(path, download);
+                if let Some(ref sender) = bytecode_compilation_sender {
+                    sender.send(installation.clone())?;
+                }
                 changelog.installed.insert(installation.key().clone());
                 for request in &requests {
                     // Take note of which installations satisfied which requests
@@ -1071,6 +1173,43 @@ fn create_bin_links(
     }
 }
 
+/// Attempt to compile the bytecode for a [`ManagedPythonInstallation`]'s stdlib
+async fn compile_stdlib_bytecode(
+    installation: &ManagedPythonInstallation,
+    concurrency: &Concurrency,
+    cache: &Cache,
+) -> Result<(usize, std::time::Duration)> {
+    let start = std::time::Instant::now();
+    let interpreter = Interpreter::query(installation.executable(false), cache)
+        .context("Couldn't locate the interpreter")?;
+
+    // Attempt to avoid accidentally bytecode compiling some other python
+    // installation's bytecode if the installed interpreter reports a weird
+    // stdlib path.
+    let interpreter_path = installation.path().canonicalize()?;
+    let stdlib_path = match interpreter.stdlib().canonicalize() {
+        Ok(path) if path.starts_with(interpreter_path) => path,
+        _ => {
+            warn_user_once!(
+                "The stdlib path for {} ({}) was not a subdirectory of its installation path. Standard library bytecode will not be compiled.",
+                installation.key(),
+                interpreter.stdlib().display()
+            );
+            return Ok((0, start.elapsed()));
+        }
+    };
+
+    let files = uv_installer::compile_tree(
+        &stdlib_path,
+        &installation.executable(false),
+        concurrency,
+        cache.root(),
+    )
+    .await
+    .with_context(|| format!("Error compiling bytecode in: {}", stdlib_path.display()))?;
+    Ok((files, start.elapsed()))
+}
+
 pub(crate) fn format_executables(
     event: &ChangeEvent,
     executables: &FxHashMap<PythonInstallationKey, FxHashSet<PathBuf>>,

crates/uv/src/lib.rs

@@ -1601,6 +1601,9 @@ async fn run(mut cli: Cli) -> Result<ExitStatus> {
             let args = settings::PythonInstallSettings::resolve(args, filesystem, environment);
             show_settings!(args);
 
+            // Initialize the cache.
+            let cache = cache.init().await?;
+
             commands::python_install(
                 &project_dir,
                 args.install_dir,
@@ -1617,6 +1620,9 @@ async fn run(mut cli: Cli) -> Result<ExitStatus> {
                 args.default,
                 globals.python_downloads,
                 cli.top_level.no_config,
+                args.compile_bytecode,
+                &globals.concurrency,
+                &cache,
                 globals.preview,
                 printer,
             )
@@ -1630,6 +1636,9 @@ async fn run(mut cli: Cli) -> Result<ExitStatus> {
             show_settings!(args);
             let upgrade = commands::PythonUpgrade::Enabled(commands::PythonUpgradeSource::Upgrade);
 
+            // Initialize the cache.
+            let cache = cache.init().await?;
+
             commands::python_install(
                 &project_dir,
                 args.install_dir,
@@ -1646,6 +1655,9 @@ async fn run(mut cli: Cli) -> Result<ExitStatus> {
                 args.default,
                 globals.python_downloads,
                 cli.top_level.no_config,
+                args.compile_bytecode,
+                &globals.concurrency,
+                &cache,
                 globals.preview,
                 printer,
             )

crates/uv/src/settings.rs

@@ -1086,6 +1086,7 @@ pub(crate) struct PythonInstallSettings {
     pub(crate) pypy_install_mirror: Option<String>,
     pub(crate) python_downloads_json_url: Option<String>,
     pub(crate) default: bool,
+    pub(crate) compile_bytecode: bool,
 }
 
 impl PythonInstallSettings {
@@ -1125,6 +1126,7 @@ impl PythonInstallSettings {
             pypy_mirror: _,
             python_downloads_json_url: _,
             default,
+            compile_bytecode,
         } = args;
 
         Self {
@@ -1144,6 +1146,12 @@ impl PythonInstallSettings {
             pypy_install_mirror,
             python_downloads_json_url,
             default,
+            compile_bytecode: flag(
+                compile_bytecode.compile_bytecode,
+                compile_bytecode.no_compile_bytecode,
+                "compile-bytecode",
+            )
+            .unwrap_or_default(),
         }
     }
 }
@@ -1162,6 +1170,7 @@ pub(crate) struct PythonUpgradeSettings {
     pub(crate) python_downloads_json_url: Option<String>,
     pub(crate) default: bool,
     pub(crate) bin: Option<bool>,
+    pub(crate) compile_bytecode: bool,
 }
 
 impl PythonUpgradeSettings {
@@ -1199,6 +1208,7 @@ impl PythonUpgradeSettings {
             pypy_mirror: _,
             reinstall,
             python_downloads_json_url: _,
+            compile_bytecode,
         } = args;
 
         Self {
@@ -1212,6 +1222,12 @@ impl PythonUpgradeSettings {
             python_downloads_json_url,
             default,
             bin,
+            compile_bytecode: flag(
+                compile_bytecode.compile_bytecode,
+                compile_bytecode.no_compile_bytecode,
+                "compile-bytecode",
+            )
+            .unwrap_or_default(),
         }
     }
 }

crates/uv/tests/it/help.rs

@@ -580,6 +580,20 @@ fn help_subsubcommand() {
               
               If multiple Python versions are requested, uv will exit with an error.
 
+    Installer options:
+          --compile-bytecode
+              Compile Python's standard library to bytecode after installation.
+              
+              By default, uv does not compile Python (`.py`) files to bytecode (`__pycache__/*.pyc`);
+              instead, compilation is performed lazily the first time a module is imported. For
+              use-cases in which start time is critical, such as CLI applications and Docker containers,
+              this option can be enabled to trade longer installation times for faster start times.
+              
+              When enabled, uv will process the Python version's `stdlib` directory. Like pip, it will
+              also ignore errors.
+              
+              [env: UV_COMPILE_BYTECODE=]
+
     Cache options:
       -n, --no-cache
               Avoid reading from or writing to the cache, instead using a temporary directory for the
@@ -834,6 +848,10 @@ fn help_flag_subsubcommand() {
           --default
               Use as the default Python version
 
+    Installer options:
+          --compile-bytecode  Compile Python's standard library to bytecode after installation [env:
+                              UV_COMPILE_BYTECODE=]
+
     Cache options:
       -n, --no-cache               Avoid reading from or writing to the cache, instead using a temporary
                                    directory for the duration of the operation [env: UV_NO_CACHE=]

crates/uv/tests/it/python_install.rs

@@ -4,6 +4,7 @@ use std::path::PathBuf;
 use std::{env, path::Path, process::Command};
 
 use crate::common::{TestContext, uv_snapshot};
+use anyhow::Context;
 use assert_cmd::assert::OutputAssertExt;
 use assert_fs::{
     assert::PathAssert,
@@ -15,6 +16,7 @@ use tracing::debug;
 
 use uv_fs::Simplified;
 use uv_static::EnvVars;
+use walkdir::WalkDir;
 
 #[test]
 fn python_install() {
@@ -3950,3 +3952,195 @@ fn python_install_upgrade_version_file() {
     hint: The version request came from a `.python-version` file; change the patch version in the file to upgrade instead
     ");
 }
+
+#[test]
+fn python_install_compile_bytecode() -> anyhow::Result<()> {
+    fn count_files_by_ext(dir: &Path, extension: &str) -> anyhow::Result<usize> {
+        let mut count = 0;
+        let walker = WalkDir::new(dir).into_iter();
+        for entry in walker {
+            let entry = entry?;
+            let path = entry.path();
+            if entry.metadata()?.is_file() && path.extension().is_some_and(|ext| ext == extension) {
+                count += 1;
+            }
+        }
+        Ok(count)
+    }
+
+    let context: TestContext = TestContext::new_with_versions(&[])
+        .with_filtered_python_keys()
+        .with_filtered_exe_suffix()
+        .with_managed_python_dirs()
+        .with_empty_python_install_mirror()
+        .with_python_download_cache();
+
+    // Install 3.14 and compile its bytecode
+    uv_snapshot!(context.filters(), context.python_install().arg("--compile-bytecode").arg("3.14"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Installed Python 3.14.2 in [TIME]
+     + cpython-3.14.2-[PLATFORM] (python3.14)
+    Bytecode compiled 1052 files in [TIME]
+    ");
+
+    // Find the stdlib path for cpython 3.14
+    let stdlib = fs_err::read_link(
+        context
+            .bin_dir
+            .child(format!("python3.14{}", std::env::consts::EXE_SUFFIX)),
+    )?
+    .parent()
+    .context("Python binary should be a child of `bin`")?
+    .parent()
+    .context("`bin` directory should be a child of the installation path")?
+    .join("lib")
+    .join("python3.14");
+
+    // And the count should match
+    let pyc_count = count_files_by_ext(&stdlib, "pyc")?;
+    let py_count = count_files_by_ext(&stdlib, "py")?;
+    assert_eq!(pyc_count, py_count);
+
+    // Attempting to install with --compile-bytecode should (currently)
+    // unconditionally re-run the bytecode compiler
+    uv_snapshot!(context.filters(), context.python_install().arg("--compile-bytecode").arg("3.14"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Python 3.14 is already installed
+    Bytecode compiled 1052 files in [TIME]
+    ");
+
+    // Reinstalling with --compile-bytecode should compile bytecode.
+    uv_snapshot!(context.filters(), context.python_install().arg("--reinstall").arg("--compile-bytecode").arg("3.14"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Installed Python 3.14.2 in [TIME]
+     ~ cpython-3.14.2-[PLATFORM] (python3.14)
+    Bytecode compiled 1052 files in [TIME]
+    ");
+
+    Ok(())
+}
+
+#[test]
+fn python_install_compile_bytecode_existing() {
+    let context: TestContext = TestContext::new_with_versions(&[])
+        .with_filtered_python_keys()
+        .with_filtered_exe_suffix()
+        .with_managed_python_dirs()
+        .with_empty_python_install_mirror()
+        .with_python_download_cache();
+
+    // A fresh install should be able to be compiled later
+    uv_snapshot!(context.filters(), context.python_install().arg("3.14"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Installed Python 3.14.2 in [TIME]
+     + cpython-3.14.2-[PLATFORM] (python3.14)
+    ");
+
+    uv_snapshot!(context.filters(), context.python_install().arg("--compile-bytecode").arg("3.14"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Python 3.14 is already installed
+    Bytecode compiled 1052 files in [TIME]
+    ");
+}
+
+#[test]
+fn python_install_compile_bytecode_upgrade() {
+    let context: TestContext = TestContext::new_with_versions(&[])
+        .with_filtered_python_keys()
+        .with_filtered_exe_suffix()
+        .with_managed_python_dirs()
+        .with_empty_python_install_mirror()
+        .with_python_download_cache();
+
+    // An upgrade should also compile bytecode
+    uv_snapshot!(context.filters(), context.python_install().arg("3.14.0"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Installed Python 3.14.0 in [TIME]
+     + cpython-3.14.0-[PLATFORM] (python3.14)
+    ");
+
+    uv_snapshot!(context.filters(), context.python_install().arg("--upgrade").arg("--compile-bytecode").arg("3.14"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Installed Python 3.14.2 in [TIME]
+     + cpython-3.14.2-[PLATFORM] (python3.14)
+    Bytecode compiled 1052 files in [TIME]
+    ");
+}
+
+#[test]
+fn python_install_compile_bytecode_multiple() {
+    let context: TestContext = TestContext::new_with_versions(&[])
+        .with_filtered_python_keys()
+        .with_filtered_exe_suffix()
+        .with_managed_python_dirs()
+        .with_empty_python_install_mirror()
+        .with_python_download_cache();
+
+    // Should handle installing and compiling multiple versions correctly
+    uv_snapshot!(context.filters(), context.python_install().arg("--compile-bytecode").arg("3.14").arg("3.12"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Installed 2 versions in [TIME]
+     + cpython-3.12.12-[PLATFORM] (python3.12)
+     + cpython-3.14.2-[PLATFORM] (python3.14)
+    Bytecode compiled 2136 files in [TIME]
+    ");
+}
+
+#[test]
+fn python_install_compile_bytecode_non_cpython() {
+    let context: TestContext = TestContext::new_with_versions(&[])
+        .with_filtered_python_keys()
+        .with_filtered_exe_suffix()
+        .with_managed_python_dirs()
+        .with_empty_python_install_mirror()
+        .with_python_download_cache();
+
+    // Should handle graalpython, pyodide and pypy gracefully
+    // Currently for pyodide this means a warning complaining about the unusual
+    // sydlib.
+    uv_snapshot!(context.filters(), context.python_install().arg("--compile-bytecode").arg("cpython-3.13.2-emscripten-wasm32-musl").arg("graalpy-3.12").arg("pypy-3.11"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    warning: The stdlib path for pyodide-3.13.2-emscripten-wasm32-musl (//lib/python3.13) was not a subdirectory of its installation path. Standard library bytecode will not be compiled.
+    Installed 3 versions in [TIME]
+     + graalpy-3.12.0-[PLATFORM] (python3.12)
+     + pypy-3.11.13-[PLATFORM] (python3.11)
+     + pyodide-3.13.2-emscripten-wasm32-musl (python3.13)
+    Bytecode compiled 2767 files in [TIME]
+    ");
+}