diff --git a/crates/puffin-cli/src/commands/reporters.rs b/crates/puffin-cli/src/commands/reporters.rs index 8540e57cd..2992648f3 100644 --- a/crates/puffin-cli/src/commands/reporters.rs +++ b/crates/puffin-cli/src/commands/reporters.rs @@ -47,6 +47,41 @@ impl puffin_resolver::Reporter for ResolverReporter { } } +#[derive(Debug)] +pub(crate) struct UnzipReporter { + progress: ProgressBar, +} + +impl From for UnzipReporter { + fn from(printer: Printer) -> Self { + let progress = ProgressBar::with_draw_target(None, printer.target()); + progress.set_message("Unzipping wheels..."); + progress.set_style( + ProgressStyle::with_template("{bar:20} [{pos}/{len}] {wide_msg:.dim}").unwrap(), + ); + Self { progress } + } +} + +impl UnzipReporter { + #[must_use] + pub(crate) fn with_length(self, length: u64) -> Self { + self.progress.set_length(length); + self + } +} + +impl puffin_installer::UnzipReporter for UnzipReporter { + fn on_unzip_progress(&self, name: &PackageName, version: &Version) { + self.progress.set_message(format!("{name}=={version}")); + self.progress.inc(1); + } + + fn on_unzip_complete(&self) { + self.progress.finish_and_clear(); + } +} + #[derive(Debug)] pub(crate) struct DownloadReporter { progress: ProgressBar, diff --git a/crates/puffin-cli/src/commands/sync.rs b/crates/puffin-cli/src/commands/sync.rs index 0658214a5..8796e8b53 100644 --- a/crates/puffin-cli/src/commands/sync.rs +++ b/crates/puffin-cli/src/commands/sync.rs @@ -16,7 +16,9 @@ use puffin_interpreter::{PythonExecutable, SitePackages}; use puffin_package::package_name::PackageName; use puffin_package::requirements::Requirements; -use crate::commands::reporters::{DownloadReporter, InstallReporter, ResolverReporter}; +use crate::commands::reporters::{ + DownloadReporter, InstallReporter, ResolverReporter, UnzipReporter, +}; use crate::commands::{elapsed, ExitStatus}; use crate::printer::Printer; @@ -161,6 +163,8 @@ pub(crate) async fn sync( resolution }; + let start = std::time::Instant::now(); + let uncached = resolution .into_files() .map(RemoteDistribution::from_file) @@ -168,8 +172,8 @@ pub(crate) async fn sync( let staging = tempfile::tempdir()?; // Download any missing distributions. - let wheels = if uncached.is_empty() { - cached + let downloads = if uncached.is_empty() { + vec![] } else { let downloader = puffin_installer::Downloader::new(&client, cache) .with_reporter(DownloadReporter::from(printer).with_length(uncached.len() as u64)); @@ -190,10 +194,41 @@ pub(crate) async fn sync( .dimmed() )?; - downloads.into_iter().chain(cached).collect::>() + downloads }; + let start = std::time::Instant::now(); + + // Unzip any downloaded distributions. + let unzips = if downloads.is_empty() { + vec![] + } else { + let unzipper = puffin_installer::Unzipper::default() + .with_reporter(UnzipReporter::from(printer).with_length(downloads.len() as u64)); + + let unzips = unzipper + .download(downloads, cache.unwrap_or(staging.path())) + .await?; + + let s = if unzips.len() == 1 { "" } else { "s" }; + writeln!( + printer, + "{}", + format!( + "Unzipped {} in {}", + format!("{} package{}", unzips.len(), s).bold(), + elapsed(start.elapsed()) + ) + .dimmed() + )?; + + unzips + }; + + let start = std::time::Instant::now(); + // Install the resolved distributions. + let wheels = unzips.into_iter().chain(cached).collect::>(); puffin_installer::Installer::new(&python) .with_reporter(InstallReporter::from(printer).with_length(wheels.len() as u64)) .install(&wheels)?; diff --git a/crates/puffin-installer/src/downloader.rs b/crates/puffin-installer/src/downloader.rs index 230770814..278e6f6c6 100644 --- a/crates/puffin-installer/src/downloader.rs +++ b/crates/puffin-installer/src/downloader.rs @@ -2,13 +2,10 @@ use std::path::Path; use anyhow::Result; use cacache::{Algorithm, Integrity}; -use rayon::iter::ParallelBridge; -use rayon::iter::ParallelIterator; use tokio::task::JoinSet; use tokio_util::compat::FuturesAsyncReadCompatExt; use tracing::debug; use url::Url; -use zip::ZipArchive; use pep440_rs::Version; use puffin_client::PypiClient; @@ -16,8 +13,6 @@ use puffin_package::package_name::PackageName; use crate::cache::WheelCache; use crate::distribution::RemoteDistribution; -use crate::vendor::CloneableSeekableReader; -use crate::LocalDistribution; pub struct Downloader<'a> { client: &'a PypiClient, @@ -49,7 +44,7 @@ impl<'a> Downloader<'a> { &'a self, wheels: &'a [RemoteDistribution], target: &'a Path, - ) -> Result> { + ) -> Result> { // Create the wheel cache subdirectory, if necessary. let wheel_cache = WheelCache::new(target); wheel_cache.init().await?; @@ -68,57 +63,29 @@ impl<'a> Downloader<'a> { } while let Some(result) = fetches.join_next().await.transpose()? { - downloads.push(result?); - } - - let mut wheels = Vec::with_capacity(downloads.len()); - - // Phase 2: Unpack the wheels into the cache. - let staging = tempfile::tempdir()?; - for download in downloads { - let remote = download.remote.clone(); - - debug!("Unpacking wheel: {}", remote.file().filename); - - // Unzip the wheel. - tokio::task::spawn_blocking({ - let target = staging.path().join(remote.id()); - move || unzip_wheel(download, &target) - }) - .await??; - - // Write the unzipped wheel to the target directory. - tokio::fs::rename( - staging.path().join(remote.id()), - wheel_cache.entry(&remote.id()), - ) - .await?; - - wheels.push(LocalDistribution::new( - remote.name().clone(), - remote.version().clone(), - wheel_cache.entry(&remote.id()), - )); + let result = result?; if let Some(reporter) = self.reporter.as_ref() { - reporter.on_download_progress(remote.name(), remote.version()); + reporter.on_download_progress(result.remote.name(), result.remote.version()); } + + downloads.push(result); } if let Some(reporter) = self.reporter.as_ref() { reporter.on_download_complete(); } - Ok(wheels) + Ok(downloads) } } #[derive(Debug, Clone)] -struct InMemoryDistribution { +pub struct InMemoryDistribution { /// The remote file from which this wheel was downloaded. - remote: RemoteDistribution, + pub(crate) remote: RemoteDistribution, /// The contents of the wheel. - buffer: Vec, + pub(crate) buffer: Vec, } /// Download a wheel to a given path. @@ -154,55 +121,10 @@ async fn fetch_wheel( Ok(InMemoryDistribution { remote, buffer }) } -/// Write a wheel into the target directory. -fn unzip_wheel(wheel: InMemoryDistribution, target: &Path) -> Result<()> { - // Read the wheel into a buffer. - let reader = std::io::Cursor::new(wheel.buffer); - let archive = ZipArchive::new(CloneableSeekableReader::new(reader))?; - - // Unzip in parallel. - (0..archive.len()) - .par_bridge() - .map(|file_number| { - let mut archive = archive.clone(); - let mut file = archive.by_index(file_number)?; - - // Determine the path of the file within the wheel. - let file_path = match file.enclosed_name() { - Some(path) => path.to_owned(), - None => return Ok(()), - }; - - // Create necessary parent directories. - let path = target.join(file_path); - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent)?; - } - - // Write the file. - let mut outfile = std::fs::File::create(&path)?; - std::io::copy(&mut file, &mut outfile)?; - - // Set permissions. - #[cfg(unix)] - { - use std::fs::Permissions; - use std::os::unix::fs::PermissionsExt; - - if let Some(mode) = file.unix_mode() { - std::fs::set_permissions(&path, Permissions::from_mode(mode))?; - } - } - - Ok(()) - }) - .collect::>() -} - pub trait Reporter: Send + Sync { /// Callback to invoke when a wheel is downloaded. fn on_download_progress(&self, name: &PackageName, version: &Version); - /// Callback to invoke when the download is complete. + /// Callback to invoke when the operation is complete. fn on_download_complete(&self); } diff --git a/crates/puffin-installer/src/lib.rs b/crates/puffin-installer/src/lib.rs index f2a252053..380b242cf 100644 --- a/crates/puffin-installer/src/lib.rs +++ b/crates/puffin-installer/src/lib.rs @@ -3,6 +3,7 @@ pub use downloader::{Downloader, Reporter as DownloadReporter}; pub use index::LocalIndex; pub use installer::{Installer, Reporter as InstallReporter}; pub use uninstall::uninstall; +pub use unzipper::{Reporter as UnzipReporter, Unzipper}; mod cache; mod distribution; @@ -10,4 +11,5 @@ mod downloader; mod index; mod installer; mod uninstall; +mod unzipper; mod vendor; diff --git a/crates/puffin-installer/src/unzipper.rs b/crates/puffin-installer/src/unzipper.rs new file mode 100644 index 000000000..810d38f40 --- /dev/null +++ b/crates/puffin-installer/src/unzipper.rs @@ -0,0 +1,134 @@ +use std::path::Path; + +use anyhow::Result; +use rayon::iter::ParallelBridge; +use rayon::iter::ParallelIterator; +use tracing::debug; +use zip::ZipArchive; + +use pep440_rs::Version; +use puffin_package::package_name::PackageName; + +use crate::cache::WheelCache; +use crate::downloader::InMemoryDistribution; +use crate::vendor::CloneableSeekableReader; +use crate::LocalDistribution; + +#[derive(Default)] +pub struct Unzipper { + reporter: Option>, +} + +impl Unzipper { + /// Set the [`Reporter`] to use for this unzipper. + #[must_use] + pub fn with_reporter(self, reporter: impl Reporter + 'static) -> Self { + Self { + reporter: Some(Box::new(reporter)), + } + } + + /// Install a set of wheels into a Python virtual environment. + pub async fn download( + &self, + downloads: Vec, + target: &Path, + ) -> Result> { + // Create the wheel cache subdirectory, if necessary. + let wheel_cache = WheelCache::new(target); + wheel_cache.init().await?; + + let staging = tempfile::tempdir()?; + + // Unpack the wheels into the cache. + let mut wheels = Vec::with_capacity(downloads.len()); + for download in downloads { + let remote = download.remote.clone(); + + debug!("Unpacking wheel: {}", remote.file().filename); + + // Unzip the wheel. + tokio::task::spawn_blocking({ + let target = staging.path().join(remote.id()); + move || unzip_wheel(download, &target) + }) + .await??; + + // Write the unzipped wheel to the target directory. + tokio::fs::rename( + staging.path().join(remote.id()), + wheel_cache.entry(&remote.id()), + ) + .await?; + + wheels.push(LocalDistribution::new( + remote.name().clone(), + remote.version().clone(), + wheel_cache.entry(&remote.id()), + )); + + if let Some(reporter) = self.reporter.as_ref() { + reporter.on_unzip_progress(remote.name(), remote.version()); + } + } + + if let Some(reporter) = self.reporter.as_ref() { + reporter.on_unzip_complete(); + } + + Ok(wheels) + } +} + +/// Write a wheel into the target directory. +fn unzip_wheel(wheel: InMemoryDistribution, target: &Path) -> Result<()> { + // Read the wheel into a buffer. + let reader = std::io::Cursor::new(wheel.buffer); + let archive = ZipArchive::new(CloneableSeekableReader::new(reader))?; + + // Unzip in parallel. + (0..archive.len()) + .par_bridge() + .map(|file_number| { + let mut archive = archive.clone(); + let mut file = archive.by_index(file_number)?; + + // Determine the path of the file within the wheel. + let file_path = match file.enclosed_name() { + Some(path) => path.to_owned(), + None => return Ok(()), + }; + + // Create necessary parent directories. + let path = target.join(file_path); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + + // Write the file. + let mut outfile = std::fs::File::create(&path)?; + std::io::copy(&mut file, &mut outfile)?; + + // Set permissions. + #[cfg(unix)] + { + use std::fs::Permissions; + use std::os::unix::fs::PermissionsExt; + + if let Some(mode) = file.unix_mode() { + std::fs::set_permissions(&path, Permissions::from_mode(mode))?; + } + } + + Ok(()) + }) + .collect::>() +} + +pub trait Reporter: Send + Sync { + /// Callback to invoke when a wheel is unzipped. + fn on_unzip_progress(&self, name: &PackageName, version: &Version); + + /// Callback to invoke when the operation is complete. + fn on_unzip_complete(&self); +} diff --git a/crates/puffin-resolver/requirements.txt b/crates/puffin-resolver/requirements.txt deleted file mode 100644 index 7482041ee..000000000 --- a/crates/puffin-resolver/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -mypy \ No newline at end of file diff --git a/crates/puffin-resolver/src/resolution.rs b/crates/puffin-resolver/src/resolution.rs index 7bc57a298..1c5bafad1 100644 --- a/crates/puffin-resolver/src/resolution.rs +++ b/crates/puffin-resolver/src/resolution.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::BTreeMap; use pep440_rs::Version; use puffin_client::File; @@ -6,11 +6,11 @@ use puffin_package::metadata::Metadata21; use puffin_package::package_name::PackageName; #[derive(Debug, Default)] -pub struct Resolution(HashMap); +pub struct Resolution(BTreeMap); impl Resolution { /// Create a new resolution from the given pinned packages. - pub(crate) fn new(packages: HashMap) -> Self { + pub(crate) fn new(packages: BTreeMap) -> Self { Self(packages) } diff --git a/crates/puffin-resolver/src/resolver.rs b/crates/puffin-resolver/src/resolver.rs index df5dfb219..49b0b5505 100644 --- a/crates/puffin-resolver/src/resolver.rs +++ b/crates/puffin-resolver/src/resolver.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashSet}; use std::str::FromStr; use anyhow::Result; @@ -87,8 +87,7 @@ impl<'a> Resolver<'a> { } // Resolve the requirements. - let mut resolution: HashMap = - HashMap::with_capacity(in_flight.len()); + let mut resolution: BTreeMap = BTreeMap::new(); while let Some(chunk) = package_stream.next().await { for result in chunk { diff --git a/crates/puffin-resolver/tests/resolver.rs b/crates/puffin-resolver/tests/resolver.rs index 982fb1d85..76d48eb41 100644 --- a/crates/puffin-resolver/tests/resolver.rs +++ b/crates/puffin-resolver/tests/resolver.rs @@ -99,7 +99,7 @@ async fn scipy() -> Result<()> { ) .await?; - assert_eq!(format!("{resolution}"), "scipy==1.11.2\nnumpy==1.25.2"); + assert_eq!(format!("{resolution}"), "numpy==1.25.2\nscipy==1.11.2"); Ok(()) }