Add --compile-bytecode to uv python install and uv python upgrade to compile the standard library

crates/uv-cli/src/lib.rs

@@ -6114,6 +6114,36 @@ 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 important, such as CLI applications and Docker containers, this
+    /// option can be enabled to trade longer installation times and some additional disk space for
+    /// faster start times.
+    ///
+    /// When enabled, uv will process the Python version's `stdlib` directory. It will ignore any
+    /// compilation errors.
+    #[arg(
+        long,
+        alias = "compile",
+        overrides_with("no_compile_bytecode"),
+        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
+    )]
+    pub no_compile_bytecode: bool,
+}
+
 #[derive(Args)]
 pub struct PythonInstallArgs {
     /// The directory to store the Python installation in.
@@ -6233,6 +6263,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 {
@@ -6293,6 +6326,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 futures::StreamExt;
-use futures::stream::FuturesUnordered;
+use anyhow::{Context, Error, Result};
+use futures::{StreamExt, join};
 use indexmap::IndexSet;
 use itertools::{Either, Itertools};
 use owo_colors::{AnsiColors, OwoColorize};
 use rustc_hash::{FxHashMap, FxHashSet};
-use tracing::{debug, trace};
+use tokio::sync::mpsc;
+use tracing::{debug, trace, warn};
 
+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,8 +29,9 @@ use uv_python::managed::{
     create_link_to_executable, python_executable_dir,
 };
 use uv_python::{
-    PythonDownloads, PythonInstallationKey, PythonInstallationMinorVersionKey, PythonRequest,
-    PythonVersionFile, VersionFileDiscoveryOptions, VersionFilePreference, VersionRequest,
+    ImplementationName, Interpreter, PythonDownloads, PythonInstallationKey,
+    PythonInstallationMinorVersionKey, PythonRequest, PythonVersionFile,
+    VersionFileDiscoveryOptions, VersionFilePreference, VersionRequest,
 };
 use uv_shell::Shell;
 use uv_trampoline_builder::{Launcher, LauncherKind};
@@ -191,6 +194,114 @@ 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 total_files = 0;
+        let mut total_elapsed = std::time::Duration::default();
+        let mut total_skipped = 0;
+        while let Some(installation) = receiver.recv().await {
+            if let Some((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;
+            } else {
+                total_skipped += 1;
+            }
+        }
+        Ok::<_, anyhow::Error>((total_files, total_elapsed, total_skipped))
+    };
+
+    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);
+
+    let (total_files, total_elapsed, total_skipped) = compiler_result?;
+    if total_files > 0 {
+        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(),
+                if total_skipped > 0 {
+                    format!(
+                        " (skipped {total_skipped} incompatible version{})",
+                        if total_skipped == 1 { "" } else { "s" }
+                    )
+                } else {
+                    String::new()
+                }
+                .dimmed()
+            )
+            .dimmed()
+        )?;
+    } else if total_skipped > 0 {
+        writeln!(
+            printer.stderr(),
+            "{}",
+            format!("No compatible versions to bytecode compile (skipped {total_skipped})")
+                .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 +544,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 +579,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 {
-        tasks.push(async {
+    let mut tasks = futures::stream::iter(&downloads)
+        .map(async |download| {
             (
                 *download,
                 download
@@ -477,8 +597,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 +613,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 +1194,53 @@ 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<Option<(usize, std::time::Duration)>> {
+    let start = std::time::Instant::now();
+
+    // Explicit matching so this heuristic is updated for future additions
+    match installation.implementation() {
+        ImplementationName::Pyodide => return Ok(None),
+        ImplementationName::GraalPy | ImplementationName::PyPy | ImplementationName::CPython => (),
+    }
+
+    let interpreter = Interpreter::query(installation.executable(false), cache)
+        .context("Couldn't locate the interpreter")?;
+
+    // Ensure the bytecode compilation occurs in the correct place, in case 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!(
+                "The stdlib path for {} ({}) is not a subdirectory of its installation path ({}).",
+                installation.key(),
+                interpreter.stdlib().display(),
+                interpreter_path.display()
+            );
+            return Ok(None);
+        }
+    };
+
+    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()))?;
+    if files == 0 {
+        return Ok(None);
+    }
+    Ok(Some((files, start.elapsed())))
+}
+
 pub(crate) fn format_executables(
     event: &ChangeEvent,
     executables: &FxHashMap<PythonInstallationKey, FxHashSet<PathBuf>>,

crates/uv/src/lib.rs

@@ -1604,6 +1604,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,
@@ -1620,6 +1623,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,
             )
@@ -1633,6 +1639,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,
@@ -1649,6 +1658,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,
             )