From f4c4b69cc7996c8c30f9143e066eed51b453ff89 Mon Sep 17 00:00:00 2001 From: Jo <10510431+j178@users.noreply.github.com> Date: Mon, 8 Jul 2024 04:01:35 +0800 Subject: [PATCH] Add progress bar when downloading python (#4840) ## Summary Resolves #4825 ## Test Plan ```sh $ cargo run -- python install --force --preview $ cargo run -- venv -p 3.12 --python-preference only-managed $ cargo run -- tool install --preview -p 3.12 --python-preference only-managed --force black ```` --------- Co-authored-by: Charlie Marsh --- Cargo.lock | 1 + crates/uv-python/Cargo.toml | 3 +- crates/uv-python/src/downloads.rs | 126 +++++++++++++++++------ crates/uv-python/src/installation.rs | 10 +- crates/uv-python/src/lib.rs | 2 +- crates/uv/src/commands/project/mod.rs | 5 +- crates/uv/src/commands/project/run.rs | 6 ++ crates/uv/src/commands/python/install.rs | 20 ++-- crates/uv/src/commands/reporters.rs | 94 ++++++++++++++--- crates/uv/src/commands/tool/install.rs | 4 + crates/uv/src/commands/tool/run.rs | 4 + crates/uv/src/commands/venv.rs | 4 + 12 files changed, 220 insertions(+), 59 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d56bc4907..5e49ca62b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4927,6 +4927,7 @@ dependencies = [ "tempfile", "test-log", "thiserror", + "tokio", "tokio-util", "tracing", "url", diff --git a/crates/uv-python/Cargo.toml b/crates/uv-python/Cargo.toml index 3a7d2e5fe..54e10215e 100644 --- a/crates/uv-python/Cargo.toml +++ b/crates/uv-python/Cargo.toml @@ -31,8 +31,8 @@ anyhow = { workspace = true } clap = { workspace = true, optional = true } configparser = { workspace = true } fs-err = { workspace = true, features = ["tokio"] } -itertools = { workspace = true } futures = { workspace = true } +itertools = { workspace = true } once_cell = { workspace = true } regex = { workspace = true } reqwest = { workspace = true } @@ -45,6 +45,7 @@ serde_json = { workspace = true } target-lexicon = { workspace = true } tempfile = { workspace = true } thiserror = { workspace = true } +tokio = { workspace = true } tokio-util = { workspace = true, features = ["compat"] } tracing = { workspace = true } url = { workspace = true } diff --git a/crates/uv-python/src/downloads.rs b/crates/uv-python/src/downloads.rs index b2a9330e0..3c02f8f6c 100644 --- a/crates/uv-python/src/downloads.rs +++ b/crates/uv-python/src/downloads.rs @@ -1,7 +1,21 @@ use std::fmt::Display; use std::io; use std::path::{Path, PathBuf}; +use std::pin::Pin; use std::str::FromStr; +use std::task::{Context, Poll}; + +use futures::TryStreamExt; +use thiserror::Error; +use tokio::io::{AsyncRead, ReadBuf}; +use tokio_util::compat::FuturesAsyncReadCompatExt; +use tracing::{debug, instrument}; +use url::Url; + +use pypi_types::{HashAlgorithm, HashDigest}; +use uv_client::WrappedReqwestError; +use uv_extract::hash::Hasher; +use uv_fs::{rename_with_retry, Simplified}; use crate::implementation::{ Error as ImplementationError, ImplementationName, LenientImplementationName, @@ -9,17 +23,6 @@ use crate::implementation::{ use crate::installation::PythonInstallationKey; use crate::platform::{self, Arch, Libc, Os}; use crate::{Interpreter, PythonRequest, PythonVersion, VersionRequest}; -use thiserror::Error; -use uv_client::WrappedReqwestError; - -use futures::TryStreamExt; - -use pypi_types::{HashAlgorithm, HashDigest}; -use tokio_util::compat::FuturesAsyncReadCompatExt; -use tracing::{debug, instrument}; -use url::Url; -use uv_extract::hash::Hasher; -use uv_fs::{rename_with_retry, Simplified}; #[derive(Error, Debug)] pub enum Error { @@ -400,11 +403,12 @@ impl ManagedPythonDownload { } /// Download and extract - #[instrument(skip(client, parent_path), fields(download = %self.key()))] + #[instrument(skip(client, parent_path, reporter), fields(download = % self.key()))] pub async fn fetch( &self, client: &uv_client::BaseClient, parent_path: &Path, + reporter: Option<&dyn Reporter>, ) -> Result { let url = Url::parse(self.url)?; let path = parent_path.join(self.key().to_string()); @@ -420,6 +424,11 @@ impl ManagedPythonDownload { // Ensure the request was successful. response.error_for_status_ref()?; + let size = response.content_length(); + let progress = reporter + .as_ref() + .map(|reporter| (reporter, reporter.on_download_start(&self.key, size))); + // Download and extract into a temporary directory. let temp_dir = tempfile::tempdir_in(parent_path).map_err(Error::DownloadDirError)?; @@ -427,28 +436,44 @@ impl ManagedPythonDownload { "Downloading {url} to temporary location {}", temp_dir.path().display() ); - let reader = response + + let stream = response .bytes_stream() .map_err(|err| io::Error::new(io::ErrorKind::Other, err)) .into_async_read(); + let mut hashers = self + .sha256 + .into_iter() + .map(|_| Hasher::from(HashAlgorithm::Sha256)) + .collect::>(); + let mut hasher = uv_extract::hash::HashReader::new(stream.compat(), &mut hashers); + debug!("Extracting {filename}"); + match progress { + Some((&reporter, progress)) => { + let mut reader = ProgressReader::new(&mut hasher, progress, reporter); + uv_extract::stream::archive(&mut reader, filename, temp_dir.path()) + .await + .map_err(|err| Error::ExtractError(filename.to_string(), err))?; + } + None => { + uv_extract::stream::archive(&mut hasher, filename, temp_dir.path()) + .await + .map_err(|err| Error::ExtractError(filename.to_string(), err))?; + } + }; + + hasher.finish().await.map_err(Error::HashExhaustion)?; + + if let Some((&reporter, progress)) = progress { + reporter.on_progress(&self.key, progress); + } + + // Check the hash if let Some(expected) = self.sha256 { - let mut hashers = [Hasher::from(HashAlgorithm::Sha256)]; - let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers); - uv_extract::stream::archive(&mut hasher, filename, temp_dir.path()) - .await - .map_err(|err| Error::ExtractError(filename.to_string(), err))?; - - hasher.finish().await.map_err(Error::HashExhaustion)?; - - let actual = hashers - .into_iter() - .map(HashDigest::from) - .next() - .unwrap() - .digest; + let actual = HashDigest::from(hashers.pop().unwrap()).digest; if !actual.eq_ignore_ascii_case(expected) { return Err(Error::HashMismatch { installation: self.key.to_string(), @@ -456,10 +481,6 @@ impl ManagedPythonDownload { actual: actual.to_string(), }); } - } else { - uv_extract::stream::archive(reader.compat(), filename, temp_dir.path()) - .await - .map_err(|err| Error::ExtractError(filename.to_string(), err))?; } // Extract the top-level directory. @@ -513,3 +534,46 @@ impl Display for ManagedPythonDownload { write!(f, "{}", self.key) } } + +pub trait Reporter: Send + Sync { + fn on_progress(&self, name: &PythonInstallationKey, id: usize); + fn on_download_start(&self, name: &PythonInstallationKey, size: Option) -> usize; + fn on_download_progress(&self, id: usize, inc: u64); + fn on_download_complete(&self); +} + +/// An asynchronous reader that reports progress as bytes are read. +struct ProgressReader<'a, R> { + reader: R, + index: usize, + reporter: &'a dyn Reporter, +} + +impl<'a, R> ProgressReader<'a, R> { + /// Create a new [`ProgressReader`] that wraps another reader. + fn new(reader: R, index: usize, reporter: &'a dyn Reporter) -> Self { + Self { + reader, + index, + reporter, + } + } +} + +impl AsyncRead for ProgressReader<'_, R> +where + R: AsyncRead + Unpin, +{ + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.as_mut().reader) + .poll_read(cx, buf) + .map_ok(|()| { + self.reporter + .on_download_progress(self.index, buf.filled().len() as u64); + }) + } +} diff --git a/crates/uv-python/src/installation.rs b/crates/uv-python/src/installation.rs index 003fddd17..fc0916a48 100644 --- a/crates/uv-python/src/installation.rs +++ b/crates/uv-python/src/installation.rs @@ -10,7 +10,7 @@ use uv_cache::Cache; use crate::discovery::{ find_best_python_installation, find_python_installation, EnvironmentPreference, PythonRequest, }; -use crate::downloads::{DownloadResult, ManagedPythonDownload, PythonDownloadRequest}; +use crate::downloads::{DownloadResult, ManagedPythonDownload, PythonDownloadRequest, Reporter}; use crate::implementation::LenientImplementationName; use crate::managed::{ManagedPythonInstallation, ManagedPythonInstallations}; use crate::platform::{Arch, Libc, Os}; @@ -82,13 +82,14 @@ impl PythonInstallation { python_fetch: PythonFetch, client_builder: &BaseClientBuilder<'a>, cache: &Cache, + reporter: Option<&dyn Reporter>, ) -> Result { let request = request.unwrap_or_default(); // Perform a fetch aggressively if managed Python is preferred if matches!(preference, PythonPreference::Managed) && python_fetch.is_automatic() { if let Some(request) = PythonDownloadRequest::try_from_request(&request) { - return Self::fetch(request.fill(), client_builder, cache).await; + return Self::fetch(request.fill(), client_builder, cache, reporter).await; } } @@ -103,7 +104,7 @@ impl PythonInstallation { { if let Some(request) = PythonDownloadRequest::try_from_request(&request) { debug!("Requested Python not found, checking for available download..."); - Self::fetch(request.fill(), client_builder, cache).await + Self::fetch(request.fill(), client_builder, cache, reporter).await } else { err } @@ -117,6 +118,7 @@ impl PythonInstallation { request: PythonDownloadRequest, client_builder: &BaseClientBuilder<'a>, cache: &Cache, + reporter: Option<&dyn Reporter>, ) -> Result { let installations = ManagedPythonInstallations::from_settings()?.init()?; let installations_dir = installations.root(); @@ -126,7 +128,7 @@ impl PythonInstallation { let client = client_builder.build(); info!("Fetching requested Python..."); - let result = download.fetch(&client, installations_dir).await?; + let result = download.fetch(&client, installations_dir, reporter).await?; let path = match result { DownloadResult::AlreadyAvailable(path) => path, diff --git a/crates/uv-python/src/lib.rs b/crates/uv-python/src/lib.rs index c5f5aed67..25f05a593 100644 --- a/crates/uv-python/src/lib.rs +++ b/crates/uv-python/src/lib.rs @@ -7,7 +7,7 @@ pub use crate::discovery::{ }; pub use crate::environment::PythonEnvironment; pub use crate::implementation::ImplementationName; -pub use crate::installation::PythonInstallation; +pub use crate::installation::{PythonInstallation, PythonInstallationKey}; pub use crate::interpreter::{Error as InterpreterError, Interpreter}; pub use crate::pointer_size::PointerSize; pub use crate::prefix::Prefix; diff --git a/crates/uv/src/commands/project/mod.rs b/crates/uv/src/commands/project/mod.rs index bd28c1597..c3047c891 100644 --- a/crates/uv/src/commands/project/mod.rs +++ b/crates/uv/src/commands/project/mod.rs @@ -25,7 +25,7 @@ use uv_resolver::{FlatIndex, OptionsBuilder, PythonRequirement, RequiresPython, use uv_types::{BuildIsolation, EmptyInstalledPackages, HashStrategy}; use crate::commands::pip::operations::Modifications; -use crate::commands::reporters::ResolverReporter; +use crate::commands::reporters::{PythonDownloadReporter, ResolverReporter}; use crate::commands::{pip, SharedState}; use crate::printer::Printer; use crate::settings::{InstallerSettingsRef, ResolverInstallerSettings, ResolverSettingsRef}; @@ -180,6 +180,8 @@ impl FoundInterpreter { .connectivity(connectivity) .native_tls(native_tls); + let reporter = PythonDownloadReporter::single(printer); + // Locate the Python interpreter to use in the environment let interpreter = PythonInstallation::find_or_fetch( python_request, @@ -188,6 +190,7 @@ impl FoundInterpreter { python_fetch, &client_builder, cache, + Some(&reporter), ) .await? .into_interpreter(); diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs index 88e49a762..2a2440abf 100644 --- a/crates/uv/src/commands/project/run.rs +++ b/crates/uv/src/commands/project/run.rs @@ -22,6 +22,7 @@ use uv_warnings::warn_user_once; use crate::commands::pip::operations::Modifications; use crate::commands::project::environment::CachedEnvironment; +use crate::commands::reporters::PythonDownloadReporter; use crate::commands::{project, ExitStatus, SharedState}; use crate::printer::Printer; use crate::settings::ResolverInstallerSettings; @@ -55,6 +56,8 @@ pub(crate) async fn run( // Initialize any shared state. let state = SharedState::default(); + let reporter = PythonDownloadReporter::single(printer); + // Determine whether the command to execute is a PEP 723 script. let script_interpreter = if let RunCommand::Python(target, _) = &command { if let Some(metadata) = uv_scripts::read_pep723_metadata(&target).await? { @@ -84,6 +87,7 @@ pub(crate) async fn run( python_fetch, &client_builder, cache, + Some(&reporter), ) .await? .into_interpreter(); @@ -214,6 +218,7 @@ pub(crate) async fn run( python_fetch, &client_builder, cache, + Some(&reporter), ) .await?; @@ -256,6 +261,7 @@ pub(crate) async fn run( python_fetch, &client_builder, cache, + Some(&reporter), ) .await? .into_interpreter() diff --git a/crates/uv/src/commands/python/install.rs b/crates/uv/src/commands/python/install.rs index 1a912703a..ba028f3f7 100644 --- a/crates/uv/src/commands/python/install.rs +++ b/crates/uv/src/commands/python/install.rs @@ -14,6 +14,7 @@ use uv_python::managed::{ManagedPythonInstallation, ManagedPythonInstallations}; use uv_python::{requests_from_version_file, PythonRequest}; use uv_warnings::warn_user_once; +use crate::commands::reporters::PythonDownloadReporter; use crate::commands::ExitStatus; use crate::printer::Printer; @@ -128,17 +129,20 @@ pub(crate) async fn install( .native_tls(native_tls) .build(); - let mut tasks = futures::stream::iter(downloads.iter()) + let reporter = PythonDownloadReporter::new(printer, downloads.len() as u64); + + let results = futures::stream::iter(downloads.iter()) .map(|download| async { - let _ = writeln!(printer.stderr(), "Downloading {}", download.key()); - let result = download.fetch(&client, installations_dir).await; + let result = download + .fetch(&client, installations_dir, Some(&reporter)) + .await; (download.python_version(), result) }) - .buffered(4); + .buffered(4) + .collect::>() + .await; - let mut results = Vec::new(); - while let Some(task) = tasks.next().await { - let (version, result) = task; + for (version, result) in results { let path = match result? { // We should only encounter already-available during concurrent installs DownloadResult::AlreadyAvailable(path) => path, @@ -151,10 +155,10 @@ pub(crate) async fn install( path } }; + // Ensure the installations have externally managed markers let installed = ManagedPythonInstallation::new(path.clone())?; installed.ensure_externally_managed()?; - results.push((version, path)); } let s = if downloads.len() == 1 { "" } else { "s" }; diff --git a/crates/uv/src/commands/reporters.rs b/crates/uv/src/commands/reporters.rs index efd6f5d8d..f34d8a296 100644 --- a/crates/uv/src/commands/reporters.rs +++ b/crates/uv/src/commands/reporters.rs @@ -11,6 +11,7 @@ use distribution_types::{ BuildableSource, CachedDist, DistributionMetadata, Name, SourceDist, VersionOrUrlRef, }; use uv_normalize::PackageName; +use uv_python::PythonInstallationKey; use crate::printer::Printer; @@ -23,9 +24,9 @@ struct ProgressReporter { #[derive(Debug)] enum ProgressMode { - // Reports top-level progress. + /// Reports top-level progress. Single, - // Reports progress of all concurrent download/build/checkout processes. + /// Reports progress of all concurrent download, build, and checkout processes. Multi { multi_progress: MultiProgress, state: Arc>, @@ -34,18 +35,18 @@ enum ProgressMode { #[derive(Default, Debug)] struct BarState { - // The number of bars that precede any download bars (i.e. build/checkout status). + /// The number of bars that precede any download bars (i.e. build/checkout status). headers: usize, - // A list of donwnload bar sizes, in descending order. + /// A list of download bar sizes, in descending order. sizes: Vec, - // A map of progress bars, by ID. + /// A map of progress bars, by ID. bars: FxHashMap, - // A monotonic counter for bar IDs. + /// A monotonic counter for bar IDs. id: usize, } impl BarState { - // Returns a unique ID for a new bar. + /// Returns a unique ID for a new progress bar. fn id(&mut self) -> usize { self.id += 1; self.id @@ -72,6 +73,7 @@ impl ProgressReporter { mode, } } + fn on_build_start(&self, source: &BuildableSource) -> usize { let ProgressMode::Multi { multi_progress, @@ -119,7 +121,7 @@ impl ProgressReporter { )); } - fn on_download_start(&self, name: &PackageName, size: Option) -> usize { + fn on_download_start(&self, name: String, size: Option) -> usize { let ProgressMode::Multi { multi_progress, state, @@ -148,10 +150,10 @@ impl ProgressReporter { .unwrap() .progress_chars("--"), ); - progress.set_message(name.to_string()); + progress.set_message(name); } else { progress.set_style(ProgressStyle::with_template("{wide_msg:.dim} ....").unwrap()); - progress.set_message(name.to_string()); + progress.set_message(name); progress.finish(); } @@ -279,7 +281,7 @@ impl uv_installer::PrepareReporter for PrepareReporter { } fn on_download_start(&self, name: &PackageName, size: Option) -> usize { - self.reporter.on_download_start(name, size) + self.reporter.on_download_start(name.to_string(), size) } fn on_download_progress(&self, id: usize, bytes: u64) { @@ -363,7 +365,7 @@ impl uv_resolver::ResolverReporter for ResolverReporter { } fn on_download_start(&self, name: &PackageName, size: Option) -> usize { - self.reporter.on_download_start(name, size) + self.reporter.on_download_start(name.to_string(), size) } fn on_download_progress(&self, id: usize, bytes: u64) { @@ -385,7 +387,7 @@ impl uv_distribution::Reporter for ResolverReporter { } fn on_download_start(&self, name: &PackageName, size: Option) -> usize { - self.reporter.on_download_start(name, size) + self.reporter.on_download_start(name.to_string(), size) } fn on_download_progress(&self, id: usize, bytes: u64) { @@ -441,6 +443,72 @@ impl uv_installer::InstallReporter for InstallReporter { } } +#[derive(Debug)] +pub(crate) struct PythonDownloadReporter { + reporter: ProgressReporter, + multiple: bool, +} + +impl PythonDownloadReporter { + /// Initialize a [`PythonDownloadReporter`] for a single Python download. + pub(crate) fn single(printer: Printer) -> Self { + Self::new(printer, 1) + } + + /// Initialize a [`PythonDownloadReporter`] for multiple Python downloads. + pub(crate) fn new(printer: Printer, length: u64) -> Self { + let multi_progress = MultiProgress::with_draw_target(printer.target()); + let root = multi_progress.add(ProgressBar::with_draw_target( + Some(length), + printer.target(), + )); + root.set_style( + ProgressStyle::with_template("{bar:20} [{pos}/{len}] {wide_msg:.dim}").unwrap(), + ); + + let reporter = ProgressReporter::new(root, multi_progress, printer); + + Self { + reporter, + multiple: length > 1, + } + } +} + +impl uv_python::downloads::Reporter for PythonDownloadReporter { + fn on_progress(&self, _name: &PythonInstallationKey, id: usize) { + self.reporter.on_download_complete(id); + + if self.multiple { + self.reporter.root.inc(1); + if self + .reporter + .root + .length() + .is_some_and(|len| self.reporter.root.position() == len) + { + self.reporter.root.finish_and_clear(); + } + } + } + + fn on_download_start(&self, name: &PythonInstallationKey, size: Option) -> usize { + if self.multiple { + self.reporter.root.set_message("Downloading Python..."); + } + self.reporter.on_download_start(name.to_string(), size) + } + + fn on_download_progress(&self, id: usize, inc: u64) { + self.reporter.on_download_progress(id, inc); + } + + fn on_download_complete(&self) { + self.reporter.root.set_message(""); + self.reporter.root.finish_and_clear(); + } +} + /// Like [`std::fmt::Display`], but with colors. trait ColorDisplay { fn to_color_string(&self) -> String; diff --git a/crates/uv/src/commands/tool/install.rs b/crates/uv/src/commands/tool/install.rs index 8c7e6634c..13b5b4923 100644 --- a/crates/uv/src/commands/tool/install.rs +++ b/crates/uv/src/commands/tool/install.rs @@ -25,6 +25,7 @@ use uv_tool::{entrypoint_paths, find_executable_directory, InstalledTools, Tool, use uv_warnings::warn_user_once; use crate::commands::project::{resolve_environment, sync_environment, update_environment}; +use crate::commands::reporters::PythonDownloadReporter; use crate::commands::tool::common::resolve_requirements; use crate::commands::{ExitStatus, SharedState}; use crate::printer::Printer; @@ -55,6 +56,8 @@ pub(crate) async fn install( .connectivity(connectivity) .native_tls(native_tls); + let reporter = PythonDownloadReporter::single(printer); + let python_request = python.as_deref().map(PythonRequest::parse); // Pre-emptively identify a Python interpreter. We need an interpreter to resolve any unnamed @@ -66,6 +69,7 @@ pub(crate) async fn install( python_fetch, &client_builder, cache, + Some(&reporter), ) .await? .into_interpreter(); diff --git a/crates/uv/src/commands/tool/run.rs b/crates/uv/src/commands/tool/run.rs index 57c96d971..524cd67bc 100644 --- a/crates/uv/src/commands/tool/run.rs +++ b/crates/uv/src/commands/tool/run.rs @@ -24,6 +24,7 @@ use uv_tool::InstalledTools; use uv_warnings::warn_user_once; use crate::commands::project::environment::CachedEnvironment; +use crate::commands::reporters::PythonDownloadReporter; use crate::commands::tool::common::resolve_requirements; use crate::commands::{ExitStatus, SharedState}; use crate::printer::Printer; @@ -154,6 +155,8 @@ async fn get_or_create_environment( .connectivity(connectivity) .native_tls(native_tls); + let reporter = PythonDownloadReporter::single(printer); + let python_request = python.map(PythonRequest::parse); // Discover an interpreter. @@ -164,6 +167,7 @@ async fn get_or_create_environment( python_fetch, &client_builder, cache, + Some(&reporter), ) .await? .into_interpreter(); diff --git a/crates/uv/src/commands/venv.rs b/crates/uv/src/commands/venv.rs index 08369288f..affeb4b32 100644 --- a/crates/uv/src/commands/venv.rs +++ b/crates/uv/src/commands/venv.rs @@ -28,6 +28,7 @@ use uv_python::{ use uv_resolver::{ExcludeNewer, FlatIndex}; use uv_types::{BuildContext, BuildIsolation, HashStrategy}; +use crate::commands::reporters::PythonDownloadReporter; use crate::commands::{pip, ExitStatus, SharedState}; use crate::printer::Printer; use crate::shell::Shell; @@ -131,6 +132,8 @@ async fn venv_impl( let client_builder_clone = client_builder.clone(); + let reporter = PythonDownloadReporter::single(printer); + let mut interpreter_request = python_request.map(PythonRequest::parse); if preview.is_enabled() && interpreter_request.is_none() { interpreter_request = request_from_version_file().await.into_diagnostic()?; @@ -144,6 +147,7 @@ async fn venv_impl( python_fetch, &client_builder, cache, + Some(&reporter), ) .await .into_diagnostic()?