From 864e4ed86084b4466a428295b76091e1a42f1d6e Mon Sep 17 00:00:00 2001 From: Bojan Serafimov Date: Thu, 11 Jan 2024 14:19:15 -0500 Subject: [PATCH] WIP --- crates/puffin-cli/src/commands/pip_sync.rs | 212 +++------------------ crates/puffin-installer/src/downloader.rs | 2 +- crates/puffin-resolver/src/finder.rs | 2 +- 3 files changed, 25 insertions(+), 191 deletions(-) diff --git a/crates/puffin-cli/src/commands/pip_sync.rs b/crates/puffin-cli/src/commands/pip_sync.rs index ac8907bbd..97cb1dbce 100644 --- a/crates/puffin-cli/src/commands/pip_sync.rs +++ b/crates/puffin-cli/src/commands/pip_sync.rs @@ -5,7 +5,7 @@ use itertools::Itertools; use owo_colors::OwoColorize; use tracing::debug; -use distribution_types::{IndexUrls, InstalledMetadata, LocalDist, LocalEditable, Name}; +use distribution_types::{IndexUrls, InstalledMetadata, LocalDist, LocalEditable, Name, CachedDist}; use install_wheel_rs::linker::LinkMode; use platform_host::Platform; use platform_tags::Tags; @@ -133,197 +133,31 @@ pub(crate) async fn pip_sync( .index_urls(index_urls.clone()) .build(); - // Resolve any registry-based requirements. - let remote = if remote.is_empty() { - Vec::new() - } else { - let start = std::time::Instant::now(); + // XXX + let wheel_finder = puffin_resolver::DistFinder::new(tags, &client, venv.interpreter()) + .with_reporter(FinderReporter::from(printer).with_length(remote.len() as u64)); + let downloader = Downloader::new(&cache, tags, &client, &build_dispatch) + .with_reporter(DownloadReporter::from(printer).with_length(remote.len() as u64)); + let wheel_finder = Box::leak(Box::new(wheel_finder)); + let downloader = Box::leak(Box::new(downloader)); - let wheel_finder = puffin_resolver::DistFinder::new(tags, &client, venv.interpreter()) - .with_reporter(FinderReporter::from(printer).with_length(remote.len() as u64)); - let resolution = wheel_finder.resolve(&remote).await?; - - let s = if resolution.len() == 1 { "" } else { "s" }; - writeln!( - printer, - "{}", - format!( - "Resolved {} in {}", - format!("{} package{}", resolution.len(), s).bold(), - elapsed(start.elapsed()) - ) - .dimmed() - )?; - - resolution.into_distributions().collect::>() - }; - - // TODO(konstin): Also check the cache whether any cached or installed dist is already known to - // have been yanked, we currently don't show this message on the second run anymore - for dist in &remote { - let Some(file) = dist.file() else { - continue; - }; - match &file.yanked { - None | Some(Yanked::Bool(false)) => {} - Some(Yanked::Bool(true)) => { - writeln!( - printer, - "{}{} {dist} is yanked. Refresh your lockfile to pin an un-yanked version.", - "warning".yellow().bold(), - ":".bold(), - )?; - } - Some(Yanked::Reason(reason)) => { - writeln!( - printer, - "{}{} {dist} is yanked (reason: \"{reason}\"). Refresh your lockfile to pin an un-yanked version.", - "warning".yellow().bold(), - ":".bold(), - )?; - } - } - } - - // Download, build, and unzip any missing distributions. - let wheels = if remote.is_empty() { - Vec::new() - } else { - let start = std::time::Instant::now(); - - let downloader = Downloader::new(&cache, tags, &client, &build_dispatch) - .with_reporter(DownloadReporter::from(printer).with_length(remote.len() as u64)); - - let wheels = downloader - .download(remote, &OnceMap::default()) - .await - .context("Failed to download distributions")?; - - let s = if wheels.len() == 1 { "" } else { "s" }; - writeln!( - printer, - "{}", - format!( - "Downloaded {} in {}", - format!("{} package{}", wheels.len(), s).bold(), - elapsed(start.elapsed()) - ) - .dimmed() - )?; - - wheels - }; - - // Remove any unnecessary packages. - if !extraneous.is_empty() || !reinstalls.is_empty() { - let start = std::time::Instant::now(); - - for dist_info in extraneous.iter().chain(reinstalls.iter()) { - let summary = puffin_installer::uninstall(dist_info).await?; - debug!( - "Uninstalled {} ({} file{}, {} director{})", - dist_info.name(), - summary.file_count, - if summary.file_count == 1 { "" } else { "s" }, - summary.dir_count, - if summary.dir_count == 1 { "y" } else { "ies" }, - ); - } - - let s = if extraneous.len() + reinstalls.len() == 1 { - "" - } else { - "s" - }; - writeln!( - printer, - "{}", - format!( - "Uninstalled {} in {}", - format!("{} package{}", extraneous.len() + reinstalls.len(), s).bold(), - elapsed(start.elapsed()) - ) - .dimmed() - )?; - } - - // Install the resolved distributions. - let wheels = wheels.into_iter().chain(local).collect::>(); - if !wheels.is_empty() { - let start = std::time::Instant::now(); - puffin_installer::Installer::new(&venv) - .with_link_mode(link_mode) - .with_reporter(InstallReporter::from(printer).with_length(wheels.len() as u64)) - .install(&wheels)?; - - let s = if wheels.len() == 1 { "" } else { "s" }; - writeln!( - printer, - "{}", - format!( - "Installed {} in {}", - format!("{} package{}", wheels.len(), s).bold(), - elapsed(start.elapsed()) - ) - .dimmed() - )?; - } - - // Report on any changes in the environment. - for event in extraneous - .into_iter() - .chain(reinstalls.into_iter()) - .map(|distribution| ChangeEvent { - dist: LocalDist::from(distribution), - kind: ChangeEventKind::Removed, - }) - .chain(wheels.into_iter().map(|distribution| ChangeEvent { - dist: LocalDist::from(distribution), - kind: ChangeEventKind::Added, - })) - .sorted_unstable_by(|a, b| { - a.dist - .name() - .cmp(b.dist.name()) - .then_with(|| a.kind.cmp(&b.kind)) - }) + let in_flight = OnceMap::default(); { - match event.kind { - ChangeEventKind::Added => { - writeln!( - printer, - " {} {}{}", - "+".green(), - event.dist.name().as_ref().white().bold(), - event.dist.installed_version().to_string().dimmed() - )?; - } - ChangeEventKind::Removed => { - writeln!( - printer, - " {} {}{}", - "-".red(), - event.dist.name().as_ref().white().bold(), - event.dist.installed_version().to_string().dimmed() - )?; - } - } + let wheel_finder = &wheel_finder; + let downloader = &downloader; + let in_flight = &in_flight; + let remote = Box::leak(Box::new(remote)); + use futures::StreamExt; + let _res: Vec = futures::stream::iter(remote) + .map(|req| async move { + let (_, dist) = wheel_finder.resolve_requirement(&req).await.unwrap(); + let wheel = downloader.get_wheel(dist, &in_flight).await.unwrap(); + wheel + }) + .buffer_unordered(50) + .collect() + .await; } - - // Validate that the environment is consistent. - if strict { - let site_packages = SitePackages::from_executable(&venv)?; - for diagnostic in site_packages.diagnostics()? { - writeln!( - printer, - "{}{} {}", - "warning".yellow().bold(), - ":".bold(), - diagnostic.message().bold() - )?; - } - } - Ok(ExitStatus::Success) } diff --git a/crates/puffin-installer/src/downloader.rs b/crates/puffin-installer/src/downloader.rs index 9e36bef71..21499ec99 100644 --- a/crates/puffin-installer/src/downloader.rs +++ b/crates/puffin-installer/src/downloader.rs @@ -34,7 +34,7 @@ pub enum Error { /// Download, build, and unzip a set of distributions. pub struct Downloader<'a, Context: BuildContext + Send + Sync> { database: DistributionDatabase<'a, Context>, - reporter: Option>, + pub reporter: Option>, } impl<'a, Context: BuildContext + Send + Sync> Downloader<'a, Context> { diff --git a/crates/puffin-resolver/src/finder.rs b/crates/puffin-resolver/src/finder.rs index 5f38e644a..3810f0d9d 100644 --- a/crates/puffin-resolver/src/finder.rs +++ b/crates/puffin-resolver/src/finder.rs @@ -46,7 +46,7 @@ impl<'a> DistFinder<'a> { /// Resolve a single pinned package, either as cached network request /// (version or no constraint) or by constructing a URL [`Dist`] from the /// specifier URL. - async fn resolve_requirement( + pub async fn resolve_requirement( &self, requirement: &Requirement, ) -> Result<(PackageName, Dist), ResolveError> {