diff --git a/crates/uv/src/commands/python/install.rs b/crates/uv/src/commands/python/install.rs index 5f295a235..02857cc03 100644 --- a/crates/uv/src/commands/python/install.rs +++ b/crates/uv/src/commands/python/install.rs @@ -6,11 +6,12 @@ use std::path::{Path, PathBuf}; use std::str::FromStr; use anyhow::{Context, Error, Result}; -use futures::StreamExt; +use futures::{StreamExt, join}; 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; @@ -198,6 +199,90 @@ pub(crate) async fn install( cache: &Cache, preview: Preview, printer: Printer, +) -> Result { + 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, + targets: Vec, + reinstall: bool, + upgrade: PythonUpgrade, + bin: Option, + registry: Option, + force: bool, + python_install_mirror: Option, + pypy_install_mirror: Option, + python_downloads_json_url: Option, + client_builder: BaseClientBuilder<'_>, + default: bool, + python_downloads: PythonDownloads, + no_config: bool, + bytecode_compilation_sender: Option>, + concurrency: &Concurrency, + preview: Preview, + printer: Printer, ) -> Result { let start = std::time::Instant::now(); @@ -440,19 +525,12 @@ pub(crate) async fn install( // For all satisfied installs, bytecode compile them now before any future // early return. - if compile_bytecode { - let mut satisfied = satisfied.clone(); - satisfied.sort(); - for installation in &satisfied { - compile_stdlib_bytecode(installation, concurrency, cache, &printer) - .await - .with_context(|| { - format!( - "Failed to bytecode-compile Python standard library for: {}", - installation.key() - ) - })?; - } + 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 @@ -514,15 +592,8 @@ pub(crate) async fn install( }; let installation = ManagedPythonInstallation::new(path, download); - if compile_bytecode { - compile_stdlib_bytecode(&installation, concurrency, cache, &printer) - .await - .with_context(|| { - format!( - "Failed to bytecode-compile Python standard library for: {}", - installation.key() - ) - })?; + if let Some(ref sender) = bytecode_compilation_sender { + sender.send(installation.clone())?; } changelog.installed.insert(installation.key().clone()); for request in &requests { @@ -1107,8 +1178,7 @@ async fn compile_stdlib_bytecode( installation: &ManagedPythonInstallation, concurrency: &Concurrency, cache: &Cache, - printer: &Printer, -) -> Result<()> { +) -> 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")?; @@ -1125,7 +1195,7 @@ async fn compile_stdlib_bytecode( installation.key(), interpreter.stdlib().display() ); - return Ok(()); + return Ok((0, start.elapsed())); } }; @@ -1137,20 +1207,7 @@ async fn compile_stdlib_bytecode( ) .await .with_context(|| format!("Error compiling bytecode in: {}", stdlib_path.display()))?; - let s = if files == 1 { "" } else { "s" }; - writeln!( - printer.stderr(), - "{}", - format!( - "Bytecode compiled {} {} {} {}", - format!("{files} file{s}").bold(), - "for".dimmed(), - installation.key().bold().dimmed(), - format!("in {}", elapsed(start.elapsed())).dimmed(), - ) - .dimmed() - )?; - Ok(()) + Ok((files, start.elapsed())) } pub(crate) fn format_executables( diff --git a/crates/uv/tests/it/python_install.rs b/crates/uv/tests/it/python_install.rs index e3adc3bd3..e8a49dc88 100644 --- a/crates/uv/tests/it/python_install.rs +++ b/crates/uv/tests/it/python_install.rs @@ -3969,9 +3969,9 @@ fn python_install_compile_bytecode() -> anyhow::Result<()> { ----- stdout ----- ----- stderr ----- - Bytecode compiled 1052 files for cpython-3.14.2-[PLATFORM] in [TIME] 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 @@ -4011,8 +4011,8 @@ fn python_install_compile_bytecode() -> anyhow::Result<()> { ----- stdout ----- ----- stderr ----- - Bytecode compiled 1052 files for cpython-3.14.2-[PLATFORM] in [TIME] Python 3.14 is already installed + Bytecode compiled 1052 files in [TIME] "); // Reinstalling with --compile-bytecode should compile bytecode. @@ -4022,9 +4022,9 @@ fn python_install_compile_bytecode() -> anyhow::Result<()> { ----- stdout ----- ----- stderr ----- - Bytecode compiled 1052 files for cpython-3.14.2-[PLATFORM] in [TIME] Installed Python 3.14.2 in [TIME] ~ cpython-3.14.2-[PLATFORM] (python3.14) + Bytecode compiled 1052 files in [TIME] "); Ok(()) @@ -4056,8 +4056,8 @@ fn python_install_compile_bytecode_existing() { ----- stdout ----- ----- stderr ----- - Bytecode compiled 1052 files for cpython-3.14.2-[PLATFORM] in [TIME] Python 3.14 is already installed + Bytecode compiled 1052 files in [TIME] "); } @@ -4087,9 +4087,9 @@ fn python_install_compile_bytecode_upgrade() { ----- stdout ----- ----- stderr ----- - Bytecode compiled 1052 files for cpython-3.14.2-[PLATFORM] in [TIME] Installed Python 3.14.2 in [TIME] + cpython-3.14.2-[PLATFORM] (python3.14) + Bytecode compiled 1052 files in [TIME] "); } @@ -4109,11 +4109,10 @@ fn python_install_compile_bytecode_multiple() { ----- stdout ----- ----- stderr ----- - Bytecode compiled 1084 files for cpython-3.12.12-[PLATFORM] in [TIME] - Bytecode compiled 1052 files for cpython-3.14.2-[PLATFORM] in [TIME] 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] "); } @@ -4136,11 +4135,10 @@ fn python_install_compile_bytecode_non_cpython() { ----- 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. - Bytecode compiled 718 files for graalpy-3.12.0-[PLATFORM] in [TIME] - Bytecode compiled 2049 files for pypy-3.11.13-[PLATFORM] in [TIME] 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] "); }