From a0294a510c9fb037025a24d665175e1ccbc16a0e Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 9 Oct 2023 23:29:09 -0400 Subject: [PATCH] Rework `puffin sync` output to summarize (#81) This also moves away from using `tracing` for user-facing logging, instead introducing a new `Printer` abstraction. Closes #66. --- Cargo.lock | 44 +++- Cargo.toml | 2 + crates/install-wheel-rs/src/lib.rs | 2 +- crates/platform-host/Cargo.toml | 2 - crates/puffin-cli/Cargo.toml | 5 +- crates/puffin-cli/src/commands/clean.rs | 10 +- crates/puffin-cli/src/commands/compile.rs | 46 ++-- crates/puffin-cli/src/commands/freeze.rs | 3 +- crates/puffin-cli/src/commands/mod.rs | 1 + crates/puffin-cli/src/commands/reporters.rs | 118 +++++++++ crates/puffin-cli/src/commands/sync.rs | 187 +++++++++----- crates/puffin-cli/src/commands/uninstall.rs | 37 ++- crates/puffin-cli/src/logging.rs | 25 +- crates/puffin-cli/src/main.rs | 21 +- crates/puffin-cli/src/printer.rs | 39 +++ crates/puffin-installer/src/cache.rs | 18 +- crates/puffin-installer/src/distribution.rs | 9 + .../src/{install.rs => downloader.rs} | 190 +++++++------- crates/puffin-installer/src/installer.rs | 63 +++++ crates/puffin-installer/src/lib.rs | 6 +- crates/puffin-installer/src/uninstall.rs | 30 +-- crates/puffin-interpreter/src/lib.rs | 2 +- crates/puffin-package/Cargo.toml | 1 - crates/puffin-package/src/requirements.rs | 4 +- crates/puffin-resolver/src/error.rs | 21 ++ crates/puffin-resolver/src/lib.rs | 231 +----------------- crates/puffin-resolver/src/resolution.rs | 48 ++++ crates/puffin-resolver/src/resolver.rs | 220 +++++++++++++++++ 28 files changed, 900 insertions(+), 485 deletions(-) create mode 100644 crates/puffin-cli/src/commands/reporters.rs create mode 100644 crates/puffin-cli/src/printer.rs rename crates/puffin-installer/src/{install.rs => downloader.rs} (50%) create mode 100644 crates/puffin-installer/src/installer.rs create mode 100644 crates/puffin-resolver/src/error.rs create mode 100644 crates/puffin-resolver/src/resolution.rs create mode 100644 crates/puffin-resolver/src/resolver.rs diff --git a/Cargo.lock b/Cargo.lock index eb6f0cd42..33c189de6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -448,6 +448,7 @@ dependencies = [ "encode_unicode", "lazy_static", "libc", + "unicode-width", "windows-sys 0.45.0", ] @@ -497,7 +498,7 @@ dependencies = [ "clap", "criterion-plot", "is-terminal", - "itertools", + "itertools 0.10.5", "num-traits", "once_cell", "oorandom", @@ -518,7 +519,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1" dependencies = [ "cast", - "itertools", + "itertools 0.10.5", ] [[package]] @@ -1122,6 +1123,19 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "indicatif" +version = "0.17.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb28741c9db9a713d93deb3bb9515c20788cef5815265bee4980e87bde7e0f25" +dependencies = [ + "console", + "instant", + "number_prefix", + "portable-atomic", + "unicode-width", +] + [[package]] name = "indoc" version = "1.0.9" @@ -1227,6 +1241,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.9" @@ -1441,6 +1464,12 @@ dependencies = [ "libc", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "object" version = "0.32.1" @@ -1635,7 +1664,6 @@ version = "0.0.1" dependencies = [ "glibc_version", "goblin", - "pep440_rs", "platform-info", "plist", "regex", @@ -1704,6 +1732,12 @@ dependencies = [ "plotters-backend", ] +[[package]] +name = "portable-atomic" +version = "1.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31114a898e107c51bb1609ffaf55a0e011cf6a4d7f1170d0015a165082c0338b" + [[package]] name = "ppv-lite86" version = "0.2.17" @@ -1739,6 +1773,9 @@ dependencies = [ "colored", "directories", "futures", + "indicatif", + "install-wheel-rs", + "itertools 0.11.0", "pep440_rs", "pep508_rs", "platform-host", @@ -1821,7 +1858,6 @@ dependencies = [ "once_cell", "pep440_rs", "pep508_rs", - "platform-host", "regex", "rfc2047-decoder", "serde", diff --git a/Cargo.toml b/Cargo.toml index ba7f5729b..acd4da17d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,8 @@ futures = { version = "0.3.28" } glibc_version = { version = "0.1.2" } goblin = { version = "0.7.1" } http-cache-reqwest = { version = "0.11.3" } +indicatif = { version = "0.17.7" } +itertools = { version = "0.11.0" } mailparse = { version = "0.14.0" } memchr = { version = "2.6.4" } once_cell = { version = "1.18.0" } diff --git a/crates/install-wheel-rs/src/lib.rs b/crates/install-wheel-rs/src/lib.rs index 6c848a0f5..0190bd84a 100644 --- a/crates/install-wheel-rs/src/lib.rs +++ b/crates/install-wheel-rs/src/lib.rs @@ -10,7 +10,7 @@ pub use install_location::{normalize_name, InstallLocation, LockedDir}; use platform_host::{Arch, Os}; pub use record::RecordEntry; pub use script::Script; -pub use uninstall::uninstall_wheel; +pub use uninstall::{uninstall_wheel, Uninstall}; pub use wheel::{ get_script_launcher, install_wheel, parse_key_value_file, read_record_file, relative_to, SHEBANG_PYTHON, diff --git a/crates/platform-host/Cargo.toml b/crates/platform-host/Cargo.toml index 439f2e808..243558950 100644 --- a/crates/platform-host/Cargo.toml +++ b/crates/platform-host/Cargo.toml @@ -10,8 +10,6 @@ authors = { workspace = true } license = { workspace = true } [dependencies] -pep440_rs = { path = "../pep440-rs" } - glibc_version = { workspace = true } goblin = { workspace = true } platform-info = { workspace = true } diff --git a/crates/puffin-cli/Cargo.toml b/crates/puffin-cli/Cargo.toml index af5b0b19e..897063289 100644 --- a/crates/puffin-cli/Cargo.toml +++ b/crates/puffin-cli/Cargo.toml @@ -8,14 +8,15 @@ name = "puffin" path = "src/main.rs" [dependencies] +install-wheel-rs = { path = "../install-wheel-rs", default-features = false } pep440_rs = { path = "../pep440-rs" } pep508_rs = { path = "../pep508-rs" } +platform-host = { path = "../platform-host" } platform-tags = { path = "../platform-tags" } puffin-client = { path = "../puffin-client" } puffin-installer = { path = "../puffin-installer" } puffin-interpreter = { path = "../puffin-interpreter" } puffin-package = { path = "../puffin-package" } -platform-host = { path = "../platform-host" } puffin-resolver = { path = "../puffin-resolver" } anyhow = { workspace = true } @@ -25,6 +26,8 @@ clap = { workspace = true, features = ["derive"] } colored = { workspace = true } directories = { workspace = true } futures = { workspace = true } +indicatif = { workspace = true } +itertools = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } diff --git a/crates/puffin-cli/src/commands/clean.rs b/crates/puffin-cli/src/commands/clean.rs index 5cd4d2268..01495e315 100644 --- a/crates/puffin-cli/src/commands/clean.rs +++ b/crates/puffin-cli/src/commands/clean.rs @@ -1,22 +1,24 @@ +use std::fmt::Write; use std::path::Path; use anyhow::{Context, Result}; -use tracing::info; +use tracing::debug; use crate::commands::ExitStatus; +use crate::printer::Printer; /// Clear the cache. -pub(crate) async fn clean(cache: Option<&Path>) -> Result { +pub(crate) async fn clean(cache: Option<&Path>, mut printer: Printer) -> Result { let Some(cache) = cache else { return Err(anyhow::anyhow!("No cache found")); }; if !cache.exists() { - info!("No cache found at: {}", cache.display()); + writeln!(printer, "No cache found at: {}", cache.display())?; return Ok(ExitStatus::Success); } - info!("Clearing cache at: {}", cache.display()); + debug!("Clearing cache at: {}", cache.display()); for entry in cache .read_dir() diff --git a/crates/puffin-cli/src/commands/compile.rs b/crates/puffin-cli/src/commands/compile.rs index 7506628a3..7d9f1dc69 100644 --- a/crates/puffin-cli/src/commands/compile.rs +++ b/crates/puffin-cli/src/commands/compile.rs @@ -1,8 +1,10 @@ +use std::fmt::Write; use std::path::Path; use std::str::FromStr; use anyhow::Result; -use tracing::{debug, info}; +use colored::Colorize; +use tracing::debug; use platform_host::Platform; use platform_tags::Tags; @@ -10,10 +12,16 @@ use puffin_client::PypiClientBuilder; use puffin_interpreter::PythonExecutable; use puffin_package::requirements::Requirements; +use crate::commands::reporters::ResolverReporter; use crate::commands::{elapsed, ExitStatus}; +use crate::printer::Printer; /// Resolve a set of requirements into a set of pinned versions. -pub(crate) async fn compile(src: &Path, cache: Option<&Path>) -> Result { +pub(crate) async fn compile( + src: &Path, + cache: Option<&Path>, + mut printer: Printer, +) -> Result { let start = std::time::Instant::now(); // Read the `requirements.txt` from disk. @@ -23,7 +31,7 @@ pub(crate) async fn compile(src: &Path, cache: Option<&Path>) -> Result) -> Result) -> Result { +pub(crate) async fn freeze(cache: Option<&Path>, _printer: Printer) -> Result { // Detect the current Python interpreter. let platform = Platform::current()?; let python = PythonExecutable::from_env(platform, cache)?; diff --git a/crates/puffin-cli/src/commands/mod.rs b/crates/puffin-cli/src/commands/mod.rs index 7fa4ffa89..2404fc5ed 100644 --- a/crates/puffin-cli/src/commands/mod.rs +++ b/crates/puffin-cli/src/commands/mod.rs @@ -10,6 +10,7 @@ pub(crate) use uninstall::uninstall; mod clean; mod compile; mod freeze; +mod reporters; mod sync; mod uninstall; diff --git a/crates/puffin-cli/src/commands/reporters.rs b/crates/puffin-cli/src/commands/reporters.rs new file mode 100644 index 000000000..8540e57cd --- /dev/null +++ b/crates/puffin-cli/src/commands/reporters.rs @@ -0,0 +1,118 @@ +use indicatif::{ProgressBar, ProgressStyle}; + +use pep440_rs::Version; +use puffin_package::package_name::PackageName; + +use crate::printer::Printer; + +#[derive(Debug)] +pub(crate) struct ResolverReporter { + progress: ProgressBar, +} + +impl From for ResolverReporter { + fn from(printer: Printer) -> Self { + let progress = ProgressBar::with_draw_target(None, printer.target()); + progress.set_message("Resolving dependencies..."); + progress.set_style( + ProgressStyle::with_template("{bar:20} [{pos}/{len}] {wide_msg:.dim}").unwrap(), + ); + Self { progress } + } +} + +impl ResolverReporter { + #[must_use] + pub(crate) fn with_length(self, length: u64) -> Self { + self.progress.set_length(length); + self + } +} + +impl puffin_resolver::Reporter for ResolverReporter { + fn on_dependency_added(&self) { + self.progress.inc_length(1); + } + + fn on_resolve_progress(&self, package: &puffin_resolver::PinnedPackage) { + self.progress.set_message(format!( + "{}=={}", + package.metadata.name, package.metadata.version + )); + self.progress.inc(1); + } + + fn on_resolve_complete(&self) { + self.progress.finish_and_clear(); + } +} + +#[derive(Debug)] +pub(crate) struct DownloadReporter { + progress: ProgressBar, +} + +impl From for DownloadReporter { + fn from(printer: Printer) -> Self { + let progress = ProgressBar::with_draw_target(None, printer.target()); + progress.set_message("Downloading wheels..."); + progress.set_style( + ProgressStyle::with_template("{bar:20} [{pos}/{len}] {wide_msg:.dim}").unwrap(), + ); + Self { progress } + } +} + +impl DownloadReporter { + #[must_use] + pub(crate) fn with_length(self, length: u64) -> Self { + self.progress.set_length(length); + self + } +} + +impl puffin_installer::DownloadReporter for DownloadReporter { + fn on_download_progress(&self, name: &PackageName, version: &Version) { + self.progress.set_message(format!("{name}=={version}")); + self.progress.inc(1); + } + + fn on_download_complete(&self) { + self.progress.finish_and_clear(); + } +} + +#[derive(Debug)] +pub(crate) struct InstallReporter { + progress: ProgressBar, +} + +impl From for InstallReporter { + fn from(printer: Printer) -> Self { + let progress = ProgressBar::with_draw_target(None, printer.target()); + progress.set_message("Installing wheels..."); + progress.set_style( + ProgressStyle::with_template("{bar:20} [{pos}/{len}] {wide_msg:.dim}").unwrap(), + ); + Self { progress } + } +} + +impl InstallReporter { + #[must_use] + pub(crate) fn with_length(self, length: u64) -> Self { + self.progress.set_length(length); + self + } +} + +impl puffin_installer::InstallReporter for InstallReporter { + fn on_install_progress(&self, name: &PackageName, version: &Version) { + self.progress.set_message(format!("{name}=={version}")); + self.progress.inc(1); + } + + fn on_install_complete(&self) { + self.progress.finish_and_clear(); + } +} diff --git a/crates/puffin-cli/src/commands/sync.rs b/crates/puffin-cli/src/commands/sync.rs index 4f8e55e54..c70cb9865 100644 --- a/crates/puffin-cli/src/commands/sync.rs +++ b/crates/puffin-cli/src/commands/sync.rs @@ -1,19 +1,24 @@ +use std::fmt::Write; use std::path::Path; use std::str::FromStr; use anyhow::Result; use bitflags::bitflags; -use tracing::{debug, info}; +use colored::Colorize; +use itertools::{Either, Itertools}; +use tracing::debug; use platform_host::Platform; use platform_tags::Tags; use puffin_client::PypiClientBuilder; -use puffin_installer::{Distribution, LocalDistribution, LocalIndex, RemoteDistribution}; +use puffin_installer::{LocalIndex, RemoteDistribution}; 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::{elapsed, ExitStatus}; +use crate::printer::Printer; bitflags! { #[derive(Debug, Copy, Clone, Default)] @@ -24,7 +29,12 @@ bitflags! { } /// Install a set of locked requirements into the current Python environment. -pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) -> Result { +pub(crate) async fn sync( + src: &Path, + cache: Option<&Path>, + flags: SyncFlags, + mut printer: Printer, +) -> Result { let start = std::time::Instant::now(); // Read the `requirements.txt` from disk. @@ -32,7 +42,10 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) -> // Parse the `requirements.txt` into a list of requirements. let requirements = Requirements::from_str(&requirements_txt)?; - let initial_requirements = requirements.len(); + if requirements.is_empty() { + writeln!(printer, "No requirements found")?; + return Ok(ExitStatus::Success); + } // Detect the current Python interpreter. let platform = Platform::current()?; @@ -60,20 +73,26 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) -> LocalIndex::default() }; - let requirements = requirements + // Filter out any already-installed or already-cached packages. + let (cached, uncached): (Vec<_>, Vec<_>) = requirements .iter() - .filter_map(|requirement| { + .filter(|requirement| { let package = PackageName::normalize(&requirement.name); // Filter out already-installed packages. if let Some(dist_info) = site_packages.get(&package) { - info!( + debug!( "Requirement already satisfied: {} ({})", package, dist_info.version() ); - return None; + false + } else { + true } + }) + .partition_map(|requirement| { + let package = PackageName::normalize(&requirement.name); // Identify any locally-available distributions that satisfy the requirement. if let Some(distribution) = local_index @@ -85,26 +104,30 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) -> distribution.name(), distribution.version() ); - return Some(Requirement::Local(distribution.clone())); + Either::Left(distribution.clone()) + } else { + debug!("Identified uncached requirement: {}", requirement); + Either::Right(requirement.clone()) } + }); - debug!("Identified uncached requirement: {}", requirement); - Some(Requirement::Remote(requirement.clone())) - }) - .collect::>(); + // Nothing to do. + if uncached.is_empty() && cached.is_empty() { + let s = if requirements.len() == 1 { "" } else { "s" }; + writeln!( + printer, + "{}", + format!( + "Audited {} in {}", + format!("{} package{}", requirements.len(), s).bold(), + elapsed(start.elapsed()) + ) + .dimmed() + )?; - if requirements.is_empty() { - let s = if initial_requirements == 1 { "" } else { "s" }; - info!( - "Audited {} package{} in {}", - initial_requirements, - s, - elapsed(start.elapsed()) - ); return Ok(ExitStatus::Success); } - // Resolve the dependencies. let client = { let mut pypi_client = PypiClientBuilder::default(); if let Some(cache) = cache { @@ -112,50 +135,90 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) -> } pypi_client.build() }; - let resolution = puffin_resolver::resolve( - requirements - .iter() - .filter_map(|requirement| match requirement { - Requirement::Remote(requirement) => Some(requirement), - Requirement::Local(_) => None, - }), - markers, - &tags, - &client, - puffin_resolver::ResolveFlags::NO_DEPS, - ) - .await?; + + // Resolve the dependencies. + let resolution = if uncached.is_empty() { + puffin_resolver::Resolution::default() + } else { + let resolver = puffin_resolver::Resolver::new(markers, &tags, &client) + .with_reporter(ResolverReporter::from(printer).with_length(uncached.len() as u64)); + let resolution = resolver + .resolve(uncached.iter(), puffin_resolver::ResolveFlags::NO_DEPS) + .await?; + + let s = if uncached.len() == 1 { "" } else { "s" }; + writeln!( + printer, + "{}", + format!( + "Resolved {} in {}", + format!("{} package{}", uncached.len(), s).bold(), + elapsed(start.elapsed()) + ) + .dimmed() + )?; + + resolution + }; + + let uncached = resolution + .into_files() + .map(RemoteDistribution::from_file) + .collect::>>()?; + let staging = tempfile::tempdir()?; + + // Download any missing distributions. + let wheels = if uncached.is_empty() { + cached + } else { + let downloader = puffin_installer::Downloader::new(&client, cache) + .with_reporter(DownloadReporter::from(printer).with_length(uncached.len() as u64)); + + let downloads = downloader + .download(&uncached, cache.unwrap_or(staging.path())) + .await?; + + let s = if uncached.len() == 1 { "" } else { "s" }; + writeln!( + printer, + "{}", + format!( + "Downloaded {} in {}", + format!("{} package{}", uncached.len(), s).bold(), + elapsed(start.elapsed()) + ) + .dimmed() + )?; + + downloads.into_iter().chain(cached).collect::>() + }; // Install the resolved distributions. - let wheels = requirements - .into_iter() - .filter_map(|requirement| match requirement { - Requirement::Remote(_) => None, - Requirement::Local(distribution) => Some(Ok(Distribution::Local(distribution))), - }) - .chain( - resolution - .into_files() - .map(|file| Ok(Distribution::Remote(RemoteDistribution::from_file(file)?))), - ) - .collect::>>()?; - puffin_installer::install(&wheels, &python, &client, cache).await?; + puffin_installer::Installer::new(&python) + .with_reporter(InstallReporter::from(printer).with_length(wheels.len() as u64)) + .install(&wheels)?; let s = if wheels.len() == 1 { "" } else { "s" }; - info!( - "Installed {} package{} in {}", - wheels.len(), - s, - elapsed(start.elapsed()) - ); + writeln!( + printer, + "{}", + format!( + "Installed {} in {}", + format!("{} package{}", wheels.len(), s).bold(), + elapsed(start.elapsed()) + ) + .dimmed() + )?; + + for wheel in wheels { + writeln!( + printer, + " {} {}{}", + "+".green(), + wheel.name().white().bold(), + format!("@{}", wheel.version()).dimmed() + )?; + } Ok(ExitStatus::Success) } - -#[derive(Debug)] -enum Requirement { - /// A requirement that must be downloaded from PyPI. - Remote(pep508_rs::Requirement), - /// A requirement that is already available locally. - Local(LocalDistribution), -} diff --git a/crates/puffin-cli/src/commands/uninstall.rs b/crates/puffin-cli/src/commands/uninstall.rs index e74283edf..bcbdc3b5b 100644 --- a/crates/puffin-cli/src/commands/uninstall.rs +++ b/crates/puffin-cli/src/commands/uninstall.rs @@ -1,15 +1,22 @@ +use std::fmt::Write; use std::path::Path; -use anyhow::Result; +use anyhow::{anyhow, Result}; use tracing::debug; use platform_host::Platform; use puffin_interpreter::PythonExecutable; +use puffin_package::package_name::PackageName; use crate::commands::ExitStatus; +use crate::printer::Printer; /// Uninstall a package from the current environment. -pub(crate) async fn uninstall(name: &str, cache: Option<&Path>) -> Result { +pub(crate) async fn uninstall( + name: &str, + cache: Option<&Path>, + mut printer: Printer, +) -> Result { // Detect the current Python interpreter. let platform = Platform::current()?; let python = PythonExecutable::from_env(platform, cache)?; @@ -18,8 +25,30 @@ pub(crate) async fn uninstall(name: &str, cache: Option<&Path>) -> Result writeln!(printer, "No files found")?, + (1, 0) => writeln!(printer, "Removed 1 file")?, + (0, 1) => writeln!(printer, "Removed 1 directory")?, + (1, 1) => writeln!(printer, "Removed 1 file and 1 directory")?, + (file_count, 0) => writeln!(printer, "Removed {file_count} files")?, + (0, dir_count) => writeln!(printer, "Removed {dir_count} directories")?, + (file_count, dir_count) => writeln!( + printer, + "Removed {file_count} files and {dir_count} directories" + )?, + } Ok(ExitStatus::Success) } diff --git a/crates/puffin-cli/src/logging.rs b/crates/puffin-cli/src/logging.rs index 19916615f..5a69e4a83 100644 --- a/crates/puffin-cli/src/logging.rs +++ b/crates/puffin-cli/src/logging.rs @@ -7,12 +7,10 @@ use tracing_tree::HierarchicalLayer; #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] pub(crate) enum Level { - /// Show deliberately user-facing messages and errors. + /// Suppress all tracing output by default (overrideable by `RUST_LOG`). #[default] Default, - /// Suppress all user-facing output. - Quiet, - /// Show all messages, including debug messages. + /// Show debug messages by default (overrideable by `RUST_LOG`). Verbose, } @@ -21,27 +19,10 @@ pub(crate) enum Level { /// /// The [`Level`] is used to dictate the default filters (which can be overridden by the `RUST_LOG` /// environment variable) along with the formatting of the output. For example, [`Level::Verbose`] -/// includes targets and timestamps, while [`Level::Default`] excludes both. +/// includes targets and timestamps, along with all `puffin=debug` messages by default. pub(crate) fn setup_logging(level: Level) { match level { Level::Default => { - // Show `INFO` messages from the CLI crate, but allow `RUST_LOG` to override. - let filter = EnvFilter::try_from_default_env() - .or_else(|_| EnvFilter::try_new("puffin=info")) - .unwrap(); - - // Regardless of the tracing level, show messages without any adornment. - tracing_subscriber::registry() - .with(filter) - .with( - tracing_subscriber::fmt::layer() - .without_time() - .with_target(false) - .with_writer(std::io::stderr), - ) - .init(); - } - Level::Quiet => { // Show nothing, but allow `RUST_LOG` to override. let filter = EnvFilter::builder() .with_default_directive(LevelFilter::OFF.into()) diff --git a/crates/puffin-cli/src/main.rs b/crates/puffin-cli/src/main.rs index 2c48d4160..1cd83e5f0 100644 --- a/crates/puffin-cli/src/main.rs +++ b/crates/puffin-cli/src/main.rs @@ -9,6 +9,7 @@ use crate::commands::ExitStatus; mod commands; mod logging; +mod printer; #[derive(Parser)] #[command(author, version, about)] @@ -85,14 +86,20 @@ struct UninstallArgs { async fn main() -> ExitCode { let cli = Cli::parse(); - logging::setup_logging(if cli.quiet { - logging::Level::Quiet - } else if cli.verbose { + logging::setup_logging(if cli.verbose { logging::Level::Verbose } else { logging::Level::Default }); + let printer = if cli.quiet { + printer::Printer::Quiet + } else if cli.verbose { + printer::Printer::Verbose + } else { + printer::Printer::Default + }; + let dirs = ProjectDirs::from("", "", "puffin"); let result = match &cli.command { @@ -102,6 +109,7 @@ async fn main() -> ExitCode { dirs.as_ref() .map(ProjectDirs::cache_dir) .filter(|_| !args.no_cache), + printer, ) .await } @@ -116,15 +124,19 @@ async fn main() -> ExitCode { } else { commands::SyncFlags::empty() }, + printer, ) .await } - Commands::Clean => commands::clean(dirs.as_ref().map(ProjectDirs::cache_dir)).await, + Commands::Clean => { + commands::clean(dirs.as_ref().map(ProjectDirs::cache_dir), printer).await + } Commands::Freeze(args) => { commands::freeze( dirs.as_ref() .map(ProjectDirs::cache_dir) .filter(|_| !args.no_cache), + printer, ) .await } @@ -134,6 +146,7 @@ async fn main() -> ExitCode { dirs.as_ref() .map(ProjectDirs::cache_dir) .filter(|_| !args.no_cache), + printer, ) .await } diff --git a/crates/puffin-cli/src/printer.rs b/crates/puffin-cli/src/printer.rs new file mode 100644 index 000000000..ce1003f6c --- /dev/null +++ b/crates/puffin-cli/src/printer.rs @@ -0,0 +1,39 @@ +use indicatif::ProgressDrawTarget; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Printer { + /// A printer that prints to standard streams (e.g., stdout). + Default, + /// A printer that suppresses all output. + Quiet, + /// A printer that prints all output, including debug messages. + Verbose, +} + +impl Printer { + pub(crate) fn target(self) -> ProgressDrawTarget { + match self { + Self::Default => ProgressDrawTarget::stderr(), + Self::Quiet => ProgressDrawTarget::hidden(), + // Confusingly, hide the progress bar when in verbose mode. + // Otherwise, it gets interleaved with debug messages. + Self::Verbose => ProgressDrawTarget::hidden(), + } + } +} + +impl std::fmt::Write for Printer { + fn write_str(&mut self, s: &str) -> std::fmt::Result { + match self { + Self::Default | Self::Verbose => { + #[allow(clippy::print_stdout)] + { + print!("{s}"); + } + } + Self::Quiet => {} + } + + Ok(()) + } +} diff --git a/crates/puffin-installer/src/cache.rs b/crates/puffin-installer/src/cache.rs index defc86ef4..2ddb79d7b 100644 --- a/crates/puffin-installer/src/cache.rs +++ b/crates/puffin-installer/src/cache.rs @@ -3,28 +3,30 @@ use std::path::{Path, PathBuf}; static WHEEL_CACHE: &str = "wheels-v0"; #[derive(Debug)] -pub(crate) struct WheelCache<'a> { - path: &'a Path, +pub(crate) struct WheelCache { + path: PathBuf, } -impl<'a> WheelCache<'a> { +impl WheelCache { /// Create a handle to the wheel cache. - pub(crate) fn new(path: &'a Path) -> Self { - Self { path } + pub(crate) fn new(root: &Path) -> Self { + Self { + path: root.join(WHEEL_CACHE), + } } /// Return the path at which a given wheel would be stored. pub(crate) fn entry(&self, id: &str) -> PathBuf { - self.path.join(WHEEL_CACHE).join(id) + self.path.join(id) } /// Initialize the wheel cache. pub(crate) async fn init(&self) -> std::io::Result<()> { - tokio::fs::create_dir_all(self.path.join(WHEEL_CACHE)).await + tokio::fs::create_dir_all(&self.path).await } /// Returns a handle to the wheel cache directory. pub(crate) async fn read_dir(&self) -> std::io::Result { - tokio::fs::read_dir(self.path.join(WHEEL_CACHE)).await + tokio::fs::read_dir(&self.path).await } } diff --git a/crates/puffin-installer/src/distribution.rs b/crates/puffin-installer/src/distribution.rs index 2ad1d31c3..e27b883b4 100644 --- a/crates/puffin-installer/src/distribution.rs +++ b/crates/puffin-installer/src/distribution.rs @@ -90,6 +90,15 @@ pub struct LocalDistribution { } impl LocalDistribution { + /// Initialize a new local distribution. + pub fn new(name: PackageName, version: Version, path: PathBuf) -> Self { + Self { + name, + version, + path, + } + } + /// Try to parse a cached distribution from a directory name (like `django-5.0a1`). pub(crate) fn try_from_path(path: &Path) -> Result> { let Some(file_name) = path.file_name() else { diff --git a/crates/puffin-installer/src/install.rs b/crates/puffin-installer/src/downloader.rs similarity index 50% rename from crates/puffin-installer/src/install.rs rename to crates/puffin-installer/src/downloader.rs index 742301d9e..230770814 100644 --- a/crates/puffin-installer/src/install.rs +++ b/crates/puffin-installer/src/downloader.rs @@ -6,123 +6,111 @@ use rayon::iter::ParallelBridge; use rayon::iter::ParallelIterator; use tokio::task::JoinSet; use tokio_util::compat::FuturesAsyncReadCompatExt; -use tracing::{debug, info}; +use tracing::debug; use url::Url; use zip::ZipArchive; +use pep440_rs::Version; use puffin_client::PypiClient; -use puffin_interpreter::PythonExecutable; +use puffin_package::package_name::PackageName; use crate::cache::WheelCache; -use crate::distribution::{Distribution, RemoteDistribution}; +use crate::distribution::RemoteDistribution; use crate::vendor::CloneableSeekableReader; +use crate::LocalDistribution; -/// Install a set of wheels into a Python virtual environment. -pub async fn install( - wheels: &[Distribution], - python: &PythonExecutable, - client: &PypiClient, - cache: Option<&Path>, -) -> Result<()> { - if wheels.is_empty() { - return Ok(()); +pub struct Downloader<'a> { + client: &'a PypiClient, + cache: Option<&'a Path>, + reporter: Option>, +} + +impl<'a> Downloader<'a> { + /// Initialize a new downloader. + pub fn new(client: &'a PypiClient, cache: Option<&'a Path>) -> Self { + Self { + client, + cache, + reporter: None, + } } - // Create the wheel cache subdirectory, if necessary. - let wheel_cache = cache.map(WheelCache::new); - if let Some(wheel_cache) = wheel_cache.as_ref() { + /// Set the [`Reporter`] to use for this downloader. + #[must_use] + pub fn with_reporter(self, reporter: impl Reporter + 'static) -> Self { + Self { + reporter: Some(Box::new(reporter)), + ..self + } + } + + /// Install a set of wheels into a Python virtual environment. + pub async fn download( + &'a self, + wheels: &'a [RemoteDistribution], + target: &'a Path, + ) -> Result> { + // Create the wheel cache subdirectory, if necessary. + let wheel_cache = WheelCache::new(target); wheel_cache.init().await?; - } - // Phase 1: Fetch the wheels in parallel. - let mut fetches = JoinSet::new(); - let mut downloads = Vec::with_capacity(wheels.len()); - for wheel in wheels { - let Distribution::Remote(remote) = wheel else { - continue; - }; + // Phase 1: Fetch the wheels in parallel. + let mut fetches = JoinSet::new(); + let mut downloads = Vec::with_capacity(wheels.len()); + for remote in wheels { + debug!("Downloading wheel: {}", remote.file().filename); - debug!("Downloading: {}", remote.file().filename); - - fetches.spawn(fetch_wheel( - remote.clone(), - client.clone(), - cache.map(Path::to_path_buf), - )); - } - - if !fetches.is_empty() { - let s = if fetches.len() == 1 { "" } else { "s" }; - info!("Downloading {} wheel{}", fetches.len(), s); - } - - while let Some(result) = fetches.join_next().await.transpose()? { - downloads.push(result?); - } - - if !downloads.is_empty() { - let s = if downloads.len() == 1 { "" } else { "s" }; - debug!("Unpacking {} wheel{}", downloads.len(), s); - } - - let staging = tempfile::tempdir()?; - - // Phase 2: Unpack the wheels into the cache. - for download in downloads { - let filename = download.remote.file().filename.clone(); - let id = download.remote.id(); - - debug!("Unpacking: {}", filename); - - // Unzip the wheel. - tokio::task::spawn_blocking({ - let target = staging.path().join(&id); - move || unzip_wheel(download, &target) - }) - .await??; - - // Write the unzipped wheel to the cache (atomically). - if let Some(wheel_cache) = wheel_cache.as_ref() { - debug!("Caching wheel: {}", filename); - tokio::fs::rename(staging.path().join(&id), wheel_cache.entry(&id)).await?; + fetches.spawn(fetch_wheel( + remote.clone(), + self.client.clone(), + self.cache.map(Path::to_path_buf), + )); } - } - let s = if wheels.len() == 1 { "" } else { "s" }; - info!( - "Linking package{}: {}", - s, - wheels - .iter() - .map(Distribution::id) - .collect::>() - .join(" ") - ); + while let Some(result) = fetches.join_next().await.transpose()? { + downloads.push(result?); + } - // Phase 3: Install each wheel. - let location = install_wheel_rs::InstallLocation::new( - python.venv().to_path_buf(), - python.simple_version(), - ); - let locked_dir = location.acquire_lock()?; + let mut wheels = Vec::with_capacity(downloads.len()); - for wheel in wheels { - match wheel { - Distribution::Remote(remote) => { - let id = remote.id(); - let dir = wheel_cache.as_ref().map_or_else( - || staging.path().join(&id), - |wheel_cache| wheel_cache.entry(&id), - ); - install_wheel_rs::unpacked::install_wheel(&locked_dir, &dir)?; - } - Distribution::Local(local) => { - install_wheel_rs::unpacked::install_wheel(&locked_dir, local.path())?; + // 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()), + )); + + if let Some(reporter) = self.reporter.as_ref() { + reporter.on_download_progress(remote.name(), remote.version()); } } - } - Ok(()) + if let Some(reporter) = self.reporter.as_ref() { + reporter.on_download_complete(); + } + + Ok(wheels) + } } #[derive(Debug, Clone)] @@ -210,3 +198,11 @@ fn unzip_wheel(wheel: InMemoryDistribution, target: &Path) -> Result<()> { }) .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. + fn on_download_complete(&self); +} diff --git a/crates/puffin-installer/src/installer.rs b/crates/puffin-installer/src/installer.rs new file mode 100644 index 000000000..01b298e7d --- /dev/null +++ b/crates/puffin-installer/src/installer.rs @@ -0,0 +1,63 @@ +use anyhow::Result; + +use pep440_rs::Version; +use puffin_interpreter::PythonExecutable; +use puffin_package::package_name::PackageName; + +use crate::LocalDistribution; + +pub struct Installer<'a> { + python: &'a PythonExecutable, + reporter: Option>, +} + +impl<'a> Installer<'a> { + /// Initialize a new installer. + pub fn new(python: &'a PythonExecutable) -> Self { + Self { + python, + reporter: None, + } + } + + /// Set the [`Reporter`] to use for this installer.. + #[must_use] + pub fn with_reporter(self, reporter: impl Reporter + 'static) -> Self { + Self { + reporter: Some(Box::new(reporter)), + ..self + } + } + + /// Install a set of wheels into a Python virtual environment. + pub fn install(self, wheels: &[LocalDistribution]) -> Result<()> { + // Install each wheel. + let location = install_wheel_rs::InstallLocation::new( + self.python.venv().to_path_buf(), + self.python.simple_version(), + ); + let locked_dir = location.acquire_lock()?; + + for wheel in wheels { + install_wheel_rs::unpacked::install_wheel(&locked_dir, wheel.path())?; + + if let Some(reporter) = self.reporter.as_ref() { + reporter.on_install_progress(wheel.name(), wheel.version()); + } + } + + if let Some(reporter) = self.reporter.as_ref() { + reporter.on_install_complete(); + } + + Ok(()) + } +} + +pub trait Reporter: Send + Sync { + /// Callback to invoke when a dependency is resolved. + fn on_install_progress(&self, name: &PackageName, version: &Version); + + /// Callback to invoke when the resolution is complete. + fn on_install_complete(&self); +} diff --git a/crates/puffin-installer/src/lib.rs b/crates/puffin-installer/src/lib.rs index c793bf765..f2a252053 100644 --- a/crates/puffin-installer/src/lib.rs +++ b/crates/puffin-installer/src/lib.rs @@ -1,11 +1,13 @@ pub use distribution::{Distribution, LocalDistribution, RemoteDistribution}; +pub use downloader::{Downloader, Reporter as DownloadReporter}; pub use index::LocalIndex; -pub use install::install; +pub use installer::{Installer, Reporter as InstallReporter}; pub use uninstall::uninstall; mod cache; mod distribution; +mod downloader; mod index; -mod install; +mod installer; mod uninstall; mod vendor; diff --git a/crates/puffin-installer/src/uninstall.rs b/crates/puffin-installer/src/uninstall.rs index 9c9624977..ae9d92c12 100644 --- a/crates/puffin-installer/src/uninstall.rs +++ b/crates/puffin-installer/src/uninstall.rs @@ -1,36 +1,14 @@ -use anyhow::{anyhow, Result}; -use tracing::info; +use anyhow::Result; -use puffin_interpreter::PythonExecutable; -use puffin_package::package_name::PackageName; +use puffin_interpreter::DistInfo; /// Uninstall a package from the specified Python environment. -pub async fn uninstall(name: &str, python: &PythonExecutable) -> Result<()> { - // Index the current `site-packages` directory. - let site_packages = puffin_interpreter::SitePackages::from_executable(python).await?; - - // Locate the package in the environment. - let Some(dist_info) = site_packages.get(&PackageName::normalize(name)) else { - return Err(anyhow!("Package not installed: {}", name)); - }; - - // Uninstall the package from the environment. +pub async fn uninstall(dist_info: &DistInfo) -> Result { let uninstall = tokio::task::spawn_blocking({ let path = dist_info.path().to_owned(); move || install_wheel_rs::uninstall_wheel(&path) }) .await??; - // Print a summary of the uninstallation. - match (uninstall.file_count, uninstall.dir_count) { - (0, 0) => info!("No files found"), - (1, 0) => info!("Removed 1 file"), - (0, 1) => info!("Removed 1 directory"), - (1, 1) => info!("Removed 1 file and 1 directory"), - (file_count, 0) => info!("Removed {file_count} files"), - (0, dir_count) => info!("Removed {dir_count} directories"), - (file_count, dir_count) => info!("Removed {file_count} files and {dir_count} directories"), - } - - Ok(()) + Ok(uninstall) } diff --git a/crates/puffin-interpreter/src/lib.rs b/crates/puffin-interpreter/src/lib.rs index e1f9eb41f..fb6a9ff91 100644 --- a/crates/puffin-interpreter/src/lib.rs +++ b/crates/puffin-interpreter/src/lib.rs @@ -7,7 +7,7 @@ use pep508_rs::MarkerEnvironment; use platform_host::Platform; use crate::python_platform::PythonPlatform; -pub use crate::site_packages::SitePackages; +pub use crate::site_packages::{DistInfo, SitePackages}; mod markers; mod python_platform; diff --git a/crates/puffin-package/Cargo.toml b/crates/puffin-package/Cargo.toml index 95c395dd3..0d53b5d31 100644 --- a/crates/puffin-package/Cargo.toml +++ b/crates/puffin-package/Cargo.toml @@ -6,7 +6,6 @@ edition = "2021" [dependencies] pep440_rs = { path = "../pep440-rs", features = ["serde"] } pep508_rs = { path = "../pep508-rs", features = ["serde"] } -platform-host = { path = "../platform-host" } anyhow = { workspace = true } mailparse = { workspace = true } diff --git a/crates/puffin-package/src/requirements.rs b/crates/puffin-package/src/requirements.rs index f2f7458f8..f3a2ca8c1 100644 --- a/crates/puffin-package/src/requirements.rs +++ b/crates/puffin-package/src/requirements.rs @@ -131,7 +131,7 @@ impl<'a> Iterator for RequirementsIterator<'a> { #[inline] fn next(&mut self) -> Option> { - if self.index == self.text.len() - 1 { + if self.index == self.text.len() { return None; } @@ -139,7 +139,7 @@ impl<'a> Iterator for RequirementsIterator<'a> { let Some((start, length)) = find_newline(&self.text[self.index..]) else { // Parse the rest of the text. let line = &self.text[self.index..]; - self.index = self.text.len() - 1; + self.index = self.text.len(); // Skip fully-commented lines. if line.trim_start().starts_with('#') { diff --git a/crates/puffin-resolver/src/error.rs b/crates/puffin-resolver/src/error.rs new file mode 100644 index 000000000..7a5f14488 --- /dev/null +++ b/crates/puffin-resolver/src/error.rs @@ -0,0 +1,21 @@ +use thiserror::Error; + +use pep508_rs::Requirement; + +#[derive(Error, Debug)] +pub enum ResolveError { + #[error("Failed to find a version of {0} that satisfies the requirement")] + NotFound(Requirement), + + #[error(transparent)] + Client(#[from] puffin_client::PypiClientError), + + #[error(transparent)] + TrySend(#[from] futures::channel::mpsc::SendError), +} + +impl From> for ResolveError { + fn from(value: futures::channel::mpsc::TrySendError) -> Self { + value.into_send_error().into() + } +} diff --git a/crates/puffin-resolver/src/lib.rs b/crates/puffin-resolver/src/lib.rs index cd627a06a..8d7efc6c5 100644 --- a/crates/puffin-resolver/src/lib.rs +++ b/crates/puffin-resolver/src/lib.rs @@ -1,227 +1,6 @@ -use std::collections::{HashMap, HashSet}; -use std::str::FromStr; +pub use resolution::{PinnedPackage, Resolution}; +pub use resolver::{Reporter, ResolveFlags, Resolver}; -use anyhow::Result; -use bitflags::bitflags; -use futures::future::Either; -use futures::{StreamExt, TryFutureExt}; -use thiserror::Error; -use tracing::{debug, info}; - -use pep440_rs::Version; -use pep508_rs::{MarkerEnvironment, Requirement}; -use platform_tags::Tags; -use puffin_client::{File, PypiClient, SimpleJson}; -use puffin_package::metadata::Metadata21; -use puffin_package::package_name::PackageName; -use wheel_filename::WheelFilename; - -#[derive(Debug, Default)] -pub struct Resolution(HashMap); - -impl Resolution { - /// Iterate over the pinned packages in this resolution. - pub fn iter(&self) -> impl Iterator { - self.0.iter() - } - - /// Iterate over the wheels in this resolution. - pub fn into_files(self) -> impl Iterator { - self.0.into_values().map(|package| package.file) - } - - /// Return the number of pinned packages in this resolution. - pub fn len(&self) -> usize { - self.0.len() - } - - /// Return `true` if there are no pinned packages in this resolution. - pub fn is_empty(&self) -> bool { - self.0.is_empty() - } -} - -#[derive(Debug)] -pub struct PinnedPackage { - metadata: Metadata21, - file: File, -} - -impl PinnedPackage { - pub fn version(&self) -> &Version { - &self.metadata.version - } -} - -bitflags! { - #[derive(Debug, Copy, Clone, Default)] - pub struct ResolveFlags: u8 { - /// Don't install package dependencies. - const NO_DEPS = 1 << 0; - } -} - -#[derive(Error, Debug)] -pub enum ResolveError { - #[error("Failed to find a version of {0} that satisfies the requirement")] - NotFound(Requirement), - - #[error(transparent)] - Client(#[from] puffin_client::PypiClientError), - - #[error(transparent)] - TrySend(#[from] futures::channel::mpsc::SendError), -} - -impl From> for ResolveError { - fn from(value: futures::channel::mpsc::TrySendError) -> Self { - value.into_send_error().into() - } -} - -/// Resolve a set of requirements into a set of pinned versions. -pub async fn resolve( - requirements: impl Iterator, - markers: &MarkerEnvironment, - tags: &Tags, - client: &PypiClient, - flags: ResolveFlags, -) -> Result { - // A channel to fetch package metadata (e.g., given `flask`, fetch all versions) and version - // metadata (e.g., given `flask==1.0.0`, fetch the metadata for that version). - let (package_sink, package_stream) = futures::channel::mpsc::unbounded(); - - // Initialize the package stream. - let mut package_stream = package_stream - .map(|request: Request| match request { - Request::Package(requirement) => Either::Left( - client - // TODO(charlie): Remove this clone. - .simple(requirement.name.clone()) - .map_ok(move |metadata| Response::Package(requirement, metadata)), - ), - Request::Version(requirement, file) => Either::Right( - client - // TODO(charlie): Remove this clone. - .file(file.clone()) - .map_ok(move |metadata| Response::Version(requirement, file, metadata)), - ), - }) - .buffer_unordered(32) - .ready_chunks(32); - - // Push all the requirements into the package sink. - let mut in_flight: HashSet = HashSet::new(); - for requirement in requirements { - debug!("--> adding root dependency: {}", requirement); - package_sink.unbounded_send(Request::Package(requirement.clone()))?; - in_flight.insert(PackageName::normalize(&requirement.name)); - } - - if in_flight.is_empty() { - return Ok(Resolution::default()); - } - - // Resolve the requirements. - let mut resolution: HashMap = - HashMap::with_capacity(in_flight.len()); - - while let Some(chunk) = package_stream.next().await { - for result in chunk { - let result: Response = result?; - match result { - Response::Package(requirement, metadata) => { - // Pick a version that satisfies the requirement. - let Some(file) = metadata.files.iter().rev().find(|file| { - // We only support wheels for now. - let Ok(name) = WheelFilename::from_str(file.filename.as_str()) else { - return false; - }; - - let Ok(version) = Version::from_str(&name.version) else { - return false; - }; - - if !name.is_compatible(tags) { - return false; - } - - requirement.is_satisfied_by(&version) - }) else { - return Err(ResolveError::NotFound(requirement)); - }; - - package_sink.unbounded_send(Request::Version(requirement, file.clone()))?; - } - Response::Version(requirement, file, metadata) => { - debug!( - "--> selected version {} for {}", - metadata.version, requirement - ); - - info!( - "Selecting: {}=={} ({})", - metadata.name, metadata.version, file.filename - ); - - // Add to the resolved set. - let normalized_name = PackageName::normalize(&requirement.name); - in_flight.remove(&normalized_name); - resolution.insert( - normalized_name, - PinnedPackage { - // TODO(charlie): Remove this clone. - metadata: metadata.clone(), - file, - }, - ); - - if !flags.intersects(ResolveFlags::NO_DEPS) { - // Enqueue its dependencies. - for dependency in metadata.requires_dist { - if !dependency.evaluate_markers( - markers, - requirement.extras.as_ref().map_or(&[], Vec::as_slice), - ) { - debug!("--> ignoring {dependency} due to environment mismatch"); - continue; - } - - let normalized_name = PackageName::normalize(&dependency.name); - - if resolution.contains_key(&normalized_name) { - continue; - } - - if !in_flight.insert(normalized_name) { - continue; - } - - debug!("--> adding transitive dependency: {}", dependency); - - package_sink.unbounded_send(Request::Package(dependency))?; - } - }; - } - } - } - - if in_flight.is_empty() { - break; - } - } - - Ok(Resolution(resolution)) -} - -#[derive(Debug)] -enum Request { - Package(Requirement), - Version(Requirement, File), -} - -#[derive(Debug)] -enum Response { - Package(Requirement, SimpleJson), - Version(Requirement, File, Metadata21), -} +mod error; +mod resolution; +mod resolver; diff --git a/crates/puffin-resolver/src/resolution.rs b/crates/puffin-resolver/src/resolution.rs new file mode 100644 index 000000000..41d3477fd --- /dev/null +++ b/crates/puffin-resolver/src/resolution.rs @@ -0,0 +1,48 @@ +use std::collections::HashMap; + +use pep440_rs::Version; +use puffin_client::File; +use puffin_package::metadata::Metadata21; +use puffin_package::package_name::PackageName; + +#[derive(Debug, Default)] +pub struct Resolution(HashMap); + +impl Resolution { + /// Create a new resolution from the given pinned packages. + pub(crate) fn new(packages: HashMap) -> Self { + Self(packages) + } + + /// Iterate over the pinned packages in this resolution. + pub fn iter(&self) -> impl Iterator { + self.0.iter() + } + + /// Iterate over the wheels in this resolution. + pub fn into_files(self) -> impl Iterator { + self.0.into_values().map(|package| package.file) + } + + /// Return the number of pinned packages in this resolution. + pub fn len(&self) -> usize { + self.0.len() + } + + /// Return `true` if there are no pinned packages in this resolution. + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } +} + +#[derive(Debug)] +pub struct PinnedPackage { + pub metadata: Metadata21, + pub file: File, +} + +impl PinnedPackage { + pub fn version(&self) -> &Version { + &self.metadata.version + } +} diff --git a/crates/puffin-resolver/src/resolver.rs b/crates/puffin-resolver/src/resolver.rs new file mode 100644 index 000000000..df5dfb219 --- /dev/null +++ b/crates/puffin-resolver/src/resolver.rs @@ -0,0 +1,220 @@ +use std::collections::{HashMap, HashSet}; +use std::str::FromStr; + +use anyhow::Result; +use bitflags::bitflags; +use futures::future::Either; +use futures::{StreamExt, TryFutureExt}; +use tracing::debug; + +use pep440_rs::Version; +use pep508_rs::{MarkerEnvironment, Requirement}; +use platform_tags::Tags; +use puffin_client::{File, PypiClient, SimpleJson}; +use puffin_package::metadata::Metadata21; +use puffin_package::package_name::PackageName; +use wheel_filename::WheelFilename; + +use crate::error::ResolveError; +use crate::resolution::{PinnedPackage, Resolution}; + +pub struct Resolver<'a> { + markers: &'a MarkerEnvironment, + tags: &'a Tags, + client: &'a PypiClient, + reporter: Option>, +} + +impl<'a> Resolver<'a> { + /// Initialize a new resolver. + pub fn new(markers: &'a MarkerEnvironment, tags: &'a Tags, client: &'a PypiClient) -> Self { + Self { + markers, + tags, + client, + reporter: None, + } + } + + /// Set the [`Reporter`] to use for this resolver. + #[must_use] + pub fn with_reporter(self, reporter: impl Reporter + 'static) -> Self { + Self { + reporter: Some(Box::new(reporter)), + ..self + } + } + + /// Resolve a set of requirements into a set of pinned versions. + pub async fn resolve( + &self, + requirements: impl Iterator, + flags: ResolveFlags, + ) -> Result { + // A channel to fetch package metadata (e.g., given `flask`, fetch all versions) and version + // metadata (e.g., given `flask==1.0.0`, fetch the metadata for that version). + let (package_sink, package_stream) = futures::channel::mpsc::unbounded(); + + // Initialize the package stream. + let mut package_stream = package_stream + .map(|request: Request| match request { + Request::Package(requirement) => Either::Left( + self.client + // TODO(charlie): Remove this clone. + .simple(requirement.name.clone()) + .map_ok(move |metadata| Response::Package(requirement, metadata)), + ), + Request::Version(requirement, file) => Either::Right( + self.client + // TODO(charlie): Remove this clone. + .file(file.clone()) + .map_ok(move |metadata| Response::Version(requirement, file, metadata)), + ), + }) + .buffer_unordered(32) + .ready_chunks(32); + + // Push all the requirements into the package sink. + let mut in_flight: HashSet = HashSet::new(); + for requirement in requirements { + debug!("Adding root dependency: {}", requirement); + package_sink.unbounded_send(Request::Package(requirement.clone()))?; + in_flight.insert(PackageName::normalize(&requirement.name)); + } + + if in_flight.is_empty() { + return Ok(Resolution::default()); + } + + // Resolve the requirements. + let mut resolution: HashMap = + HashMap::with_capacity(in_flight.len()); + + while let Some(chunk) = package_stream.next().await { + for result in chunk { + let result: Response = result?; + match result { + Response::Package(requirement, metadata) => { + // Pick a version that satisfies the requirement. + let Some(file) = metadata.files.iter().rev().find(|file| { + // We only support wheels for now. + let Ok(name) = WheelFilename::from_str(file.filename.as_str()) else { + return false; + }; + + let Ok(version) = Version::from_str(&name.version) else { + return false; + }; + + if !name.is_compatible(self.tags) { + return false; + } + + requirement.is_satisfied_by(&version) + }) else { + return Err(ResolveError::NotFound(requirement)); + }; + + package_sink.unbounded_send(Request::Version(requirement, file.clone()))?; + } + Response::Version(requirement, file, metadata) => { + debug!( + "Selecting: {}=={} ({})", + metadata.name, metadata.version, file.filename + ); + + let package = PinnedPackage { + metadata: metadata.clone(), + file, + }; + + if let Some(reporter) = self.reporter.as_ref() { + reporter.on_resolve_progress(&package); + } + + // Add to the resolved set. + let normalized_name = PackageName::normalize(&requirement.name); + in_flight.remove(&normalized_name); + resolution.insert(normalized_name, package); + + if !flags.intersects(ResolveFlags::NO_DEPS) { + // Enqueue its dependencies. + for dependency in metadata.requires_dist { + if !dependency.evaluate_markers( + self.markers, + requirement.extras.as_ref().map_or(&[], Vec::as_slice), + ) { + debug!("Ignoring {dependency} due to environment mismatch"); + continue; + } + + let normalized_name = PackageName::normalize(&dependency.name); + + if resolution.contains_key(&normalized_name) { + continue; + } + + if !in_flight.insert(normalized_name) { + continue; + } + + debug!("Adding transitive dependency: {}", dependency); + + package_sink.unbounded_send(Request::Package(dependency))?; + + if let Some(reporter) = self.reporter.as_ref() { + reporter.on_dependency_added(); + } + } + }; + } + } + } + + if in_flight.is_empty() { + break; + } + } + + if let Some(reporter) = self.reporter.as_ref() { + reporter.on_resolve_complete(); + } + + Ok(Resolution::new(resolution)) + } +} + +#[derive(Debug)] +enum Request { + /// A request to fetch the metadata for a package. + Package(Requirement), + /// A request to fetch the metadata for a specific version of a package. + Version(Requirement, File), +} + +#[derive(Debug)] +enum Response { + /// The returned metadata for a package. + Package(Requirement, SimpleJson), + /// The returned metadata for a specific version of a package. + Version(Requirement, File, Metadata21), +} + +pub trait Reporter: Send + Sync { + /// Callback to invoke when a dependency is added to the resolution. + fn on_dependency_added(&self); + + /// Callback to invoke when a dependency is resolved. + fn on_resolve_progress(&self, package: &PinnedPackage); + + /// Callback to invoke when the resolution is complete. + fn on_resolve_complete(&self); +} + +bitflags! { + #[derive(Debug, Copy, Clone, Default)] + pub struct ResolveFlags: u8 { + /// Don't install package dependencies. + const NO_DEPS = 1 << 0; + } +}