mirrors/uv
1
0
Fork 0
mirror of https://github.com/astral-sh/uv.git synced 2025-07-07 13:25:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Rework puffin sync output to summarize (#81)

Browse source 
This also moves away from using `tracing` for user-facing logging,
instead introducing a new `Printer` abstraction.

Closes #66.
This commit is contained in:
Charlie Marsh 2023-10-09 23:29:09 -04:00 • committed by GitHub
parent 2d4a8c361b
commit a0294a510c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
28 changed files with 900 additions and 485 deletions
Download patch file Download diff file Expand all files Collapse all files

44
Cargo.lock generated
View file

@ -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",

2
Cargo.toml
View file

@ -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" }

2
crates/install-wheel-rs/src/lib.rs
View file

@ -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,

2
crates/platform-host/Cargo.toml
View file

@ -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 }

5
crates/puffin-cli/Cargo.toml
View file

@ -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 }

10
crates/puffin-cli/src/commands/clean.rs
View file

@ -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<ExitStatus> {
pub(crate) async fn clean(cache: Option<&Path>, mut printer: Printer) -> Result<ExitStatus> {
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()

46
crates/puffin-cli/src/commands/compile.rs
View file

@ -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<ExitStatus> {
pub(crate) async fn compile(
src: &Path,
cache: Option<&Path>,
mut printer: Printer,
) -> Result<ExitStatus> {
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<ExitStat
let requirements = Requirements::from_str(&requirements_txt)?;
if requirements.is_empty() {
info!("No requirements found");
writeln!(printer, "No requirements found")?;
return Ok(ExitStatus::Success);
}
@ -51,22 +59,26 @@ pub(crate) async fn compile(src: &Path, cache: Option<&Path>) -> Result<ExitStat
};
// Resolve the dependencies.
let resolution = puffin_resolver::resolve(
requirements.iter(),
markers,
&tags,
&client,
puffin_resolver::ResolveFlags::default(),
)
.await?;
let resolver = puffin_resolver::Resolver::new(markers, &tags, &client)
.with_reporter(ResolverReporter::from(printer));
let resolution = resolver
.resolve(
requirements.iter(),
puffin_resolver::ResolveFlags::default(),
)
.await?;
let s = if resolution.len() == 1 { "" } else { "s" };
info!(
"Resolved {} package{} in {}",
resolution.len(),
s,
elapsed(start.elapsed())
);
writeln!(
printer,
"{}",
format!(
"Resolved {} in {}",
format!("{} package{}", resolution.len(), s).bold(),
elapsed(start.elapsed())
)
.dimmed()
)?;
for (name, package) in resolution.iter() {
#[allow(clippy::print_stdout)]

3
crates/puffin-cli/src/commands/freeze.rs
View file

@ -7,9 +7,10 @@ use platform_host::Platform;
use puffin_interpreter::{PythonExecutable, SitePackages};
use crate::commands::ExitStatus;
use crate::printer::Printer;
/// Enumerate the installed packages in the current environment.
pub(crate) async fn freeze(cache: Option<&Path>) -> Result<ExitStatus> {
pub(crate) async fn freeze(cache: Option<&Path>, _printer: Printer) -> Result<ExitStatus> {
// Detect the current Python interpreter.
let platform = Platform::current()?;
let python = PythonExecutable::from_env(platform, cache)?;

1
crates/puffin-cli/src/commands/mod.rs
View file

@ -10,6 +10,7 @@ pub(crate) use uninstall::uninstall;
mod clean;
mod compile;
mod freeze;
mod reporters;
mod sync;
mod uninstall;

118
crates/puffin-cli/src/commands/reporters.rs Normal file
View file

@ -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<Printer> 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<Printer> 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<Printer> 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();
}
}

187
crates/puffin-cli/src/commands/sync.rs
View file

@ -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<ExitStatus> {
pub(crate) async fn sync(
src: &Path,
cache: Option<&Path>,
flags: SyncFlags,
mut printer: Printer,
) -> Result<ExitStatus> {
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::<Vec<_>>();
// 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::<Result<Vec<_>>>()?;
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::<Vec<_>>()
};
// 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::<Result<Vec<_>>>()?;
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),
}

37
crates/puffin-cli/src/commands/uninstall.rs
View file

@ -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<ExitStatus> {
pub(crate) async fn uninstall(
name: &str,
cache: Option<&Path>,
mut printer: Printer,
) -> Result<ExitStatus> {
// 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<ExitSt
python.executable().display()
);
// Uninstall the package from the current environment.
puffin_installer::uninstall(name, &python).await?;
// 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.
let uninstall = puffin_installer::uninstall(dist_info).await?;
// Print a summary of the uninstallation.
match (uninstall.file_count, uninstall.dir_count) {
(0, 0) => 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)
}

25
crates/puffin-cli/src/logging.rs
View file

@ -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())

21
crates/puffin-cli/src/main.rs
View file

@ -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
}

39
crates/puffin-cli/src/printer.rs Normal file
View file

@ -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(())
}
}

18
crates/puffin-installer/src/cache.rs
View file

@ -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::ReadDir> {
tokio::fs::read_dir(self.path.join(WHEEL_CACHE)).await
tokio::fs::read_dir(&self.path).await
}
}

9
crates/puffin-installer/src/distribution.rs
View file

@ -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<Option<Self>> {
let Some(file_name) = path.file_name() else {

190
crates/puffin-installer/src/install.rs → crates/puffin-installer/src/downloader.rs
View file

@ -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<Box<dyn Reporter>>,
}
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<Vec<LocalDistribution>> {
// 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::<Vec<_>>()
.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::<Result<_>>()
}
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);
}

63
crates/puffin-installer/src/installer.rs Normal file
View file

@ -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<Box<dyn Reporter>>,
}
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);
}

6
crates/puffin-installer/src/lib.rs
View file

@ -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;

30
crates/puffin-installer/src/uninstall.rs
View file

@ -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<install_wheel_rs::Uninstall> {
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)
}

2
crates/puffin-interpreter/src/lib.rs
View file

@ -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;

1
crates/puffin-package/Cargo.toml
View file

@ -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 }

4
crates/puffin-package/src/requirements.rs
View file

@ -131,7 +131,7 @@ impl<'a> Iterator for RequirementsIterator<'a> {
#[inline]
fn next(&mut self) -> Option<RequirementLine<'a>> {
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('#') {

21
crates/puffin-resolver/src/error.rs Normal file
View file

@ -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<T> From<futures::channel::mpsc::TrySendError<T>> for ResolveError {
fn from(value: futures::channel::mpsc::TrySendError<T>) -> Self {
value.into_send_error().into()
}
}

231
crates/puffin-resolver/src/lib.rs
View file

@ -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<PackageName, PinnedPackage>);
impl Resolution {
/// Iterate over the pinned packages in this resolution.
pub fn iter(&self) -> impl Iterator<Item = (&PackageName, &PinnedPackage)> {
self.0.iter()
}
/// Iterate over the wheels in this resolution.
pub fn into_files(self) -> impl Iterator<Item = File> {
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<T> From<futures::channel::mpsc::TrySendError<T>> for ResolveError {
fn from(value: futures::channel::mpsc::TrySendError<T>) -> Self {
value.into_send_error().into()
}
}
/// Resolve a set of requirements into a set of pinned versions.
pub async fn resolve(
requirements: impl Iterator<Item = &Requirement>,
markers: &MarkerEnvironment,
tags: &Tags,
client: &PypiClient,
flags: ResolveFlags,
) -> Result<Resolution, ResolveError> {
// 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<PackageName> = 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<PackageName, PinnedPackage> =
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;

48
crates/puffin-resolver/src/resolution.rs Normal file
View file

@ -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<PackageName, PinnedPackage>);
impl Resolution {
/// Create a new resolution from the given pinned packages.
pub(crate) fn new(packages: HashMap<PackageName, PinnedPackage>) -> Self {
Self(packages)
}
/// Iterate over the pinned packages in this resolution.
pub fn iter(&self) -> impl Iterator<Item = (&PackageName, &PinnedPackage)> {
self.0.iter()
}
/// Iterate over the wheels in this resolution.
pub fn into_files(self) -> impl Iterator<Item = File> {
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
}
}

220
crates/puffin-resolver/src/resolver.rs Normal file
View file

@ -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<Box<dyn Reporter>>,
}
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<Item = &Requirement>,
flags: ResolveFlags,
) -> Result<Resolution, ResolveError> {
// 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<PackageName> = 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<PackageName, PinnedPackage> =
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;
}
}