From 82ff136a74acf729c832d7d8cdc42b0e5c499a53 Mon Sep 17 00:00:00 2001 From: konsti Date: Mon, 15 Jan 2024 04:04:55 +0100 Subject: [PATCH] Add find links supports to pip-sync (#914) Closes #877 --- .../src/prioritized_distribution.rs | 10 +++ crates/puffin-cli/src/commands/pip_sync.rs | 10 ++- crates/puffin-cli/tests/pip_sync.rs | 53 +++++++++++ crates/puffin-dev/src/install_many.rs | 9 +- crates/puffin-resolver/src/finder.rs | 90 ++++++++++++------- 5 files changed, 135 insertions(+), 37 deletions(-) diff --git a/crates/distribution-types/src/prioritized_distribution.rs b/crates/distribution-types/src/prioritized_distribution.rs index c2f3174c7..c704391de 100644 --- a/crates/distribution-types/src/prioritized_distribution.rs +++ b/crates/distribution-types/src/prioritized_distribution.rs @@ -160,6 +160,16 @@ impl PrioritizedDistribution { } } + /// Return the source distribution, if any. + pub fn source(&self) -> Option<&DistRequiresPython> { + self.source.as_ref() + } + + /// Return the compatible built distribution, if any. + pub fn compatible_wheel(&self) -> Option<&(DistRequiresPython, TagPriority)> { + self.compatible_wheel.as_ref() + } + /// Return the hashes for each distribution. pub fn hashes(&self) -> &[Hashes] { &self.hashes diff --git a/crates/puffin-cli/src/commands/pip_sync.rs b/crates/puffin-cli/src/commands/pip_sync.rs index bd0f8b400..0eb1d86eb 100644 --- a/crates/puffin-cli/src/commands/pip_sync.rs +++ b/crates/puffin-cli/src/commands/pip_sync.rs @@ -10,7 +10,7 @@ use install_wheel_rs::linker::LinkMode; use platform_host::Platform; use platform_tags::Tags; use puffin_cache::Cache; -use puffin_client::{RegistryClient, RegistryClientBuilder}; +use puffin_client::{FlatIndex, RegistryClient, RegistryClientBuilder}; use puffin_dispatch::BuildDispatch; use puffin_installer::{Downloader, InstallPlan, Reinstall, ResolvedEditable, SitePackages}; use puffin_interpreter::Virtualenv; @@ -139,8 +139,12 @@ pub(crate) async fn pip_sync( } else { let start = std::time::Instant::now(); - let wheel_finder = puffin_resolver::DistFinder::new(tags, &client, venv.interpreter()) - .with_reporter(FinderReporter::from(printer).with_length(remote.len() as u64)); + let flat_index_files = client.flat_index().await?; + let flat_index = FlatIndex::from_files(flat_index_files, tags); + + let wheel_finder = + puffin_resolver::DistFinder::new(tags, &client, venv.interpreter(), &flat_index) + .with_reporter(FinderReporter::from(printer).with_length(remote.len() as u64)); let resolution = wheel_finder.resolve(&remote).await?; let s = if resolution.len() == 1 { "" } else { "s" }; diff --git a/crates/puffin-cli/tests/pip_sync.rs b/crates/puffin-cli/tests/pip_sync.rs index ad9e85619..c2c91b1f5 100644 --- a/crates/puffin-cli/tests/pip_sync.rs +++ b/crates/puffin-cli/tests/pip_sync.rs @@ -7,6 +7,8 @@ use std::process::Command; use anyhow::{Context, Result}; use assert_cmd::prelude::*; use assert_fs::prelude::*; +use assert_fs::TempDir; +use indoc::indoc; use insta_cmd::_macro_support::insta; use insta_cmd::{assert_cmd_snapshot, get_cargo_bin}; @@ -2664,3 +2666,54 @@ fn sync_legacy_sdist_setuptools() -> Result<()> { Ok(()) } + +/// Sync using `--find-links` with a local directory. +#[test] +fn find_links() -> Result<()> { + let temp_dir = TempDir::new()?; + let cache_dir = TempDir::new()?; + let venv = create_venv_py312(&temp_dir, &cache_dir); + + let requirements_txt = temp_dir.child("requirements.txt"); + requirements_txt.write_str(indoc! {r" + markupsafe==2.1.3 + numpy==1.26.3 + tqdm==1000.0.0 + werkzeug @ https://files.pythonhosted.org/packages/c3/fc/254c3e9b5feb89ff5b9076a23218dafbc99c96ac5941e900b71206e6313b/werkzeug-3.0.1-py3-none-any.whl + "})?; + + let project_root = fs_err::canonicalize(std::env::current_dir()?.join("../.."))?; + let project_root_string = project_root.display().to_string(); + let filters: Vec<_> = iter::once((project_root_string.as_str(), "[PROJECT_ROOT]")) + .chain(INSTA_FILTERS.to_vec()) + .collect(); + + insta::with_settings!({ + filters => filters + }, { + assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME)) + .arg("pip-sync") + .arg("requirements.txt") + .arg("--find-links") + .arg(project_root.join("scripts/wheels/")) + .arg("--cache-dir") + .arg(cache_dir.path()) + .env("VIRTUAL_ENV", venv.as_os_str()) + .current_dir(&temp_dir), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Resolved 4 packages in [TIME] + Downloaded 4 packages in [TIME] + Installed 4 packages in [TIME] + + markupsafe==2.1.3 + + numpy==1.26.3 + + tqdm==1000.0.0 + + werkzeug==3.0.1 (from https://files.pythonhosted.org/packages/c3/fc/254c3e9b5feb89ff5b9076a23218dafbc99c96ac5941e900b71206e6313b/werkzeug-3.0.1-py3-none-any.whl) + "###); + }); + + Ok(()) +} diff --git a/crates/puffin-dev/src/install_many.rs b/crates/puffin-dev/src/install_many.rs index 6a62f0381..738bd90c5 100644 --- a/crates/puffin-dev/src/install_many.rs +++ b/crates/puffin-dev/src/install_many.rs @@ -103,10 +103,11 @@ async fn install_chunk( venv: &Virtualenv, index_locations: &IndexLocations, ) -> Result<()> { - let resolution: Vec<_> = DistFinder::new(tags, client, venv.interpreter()) - .resolve_stream(requirements) - .collect() - .await; + let resolution: Vec<_> = + DistFinder::new(tags, client, venv.interpreter(), &FxHashMap::default()) + .resolve_stream(requirements) + .collect() + .await; let (resolution, failures): (FxHashMap, Vec<_>) = resolution.into_iter().partition_result(); for failure in &failures { diff --git a/crates/puffin-resolver/src/finder.rs b/crates/puffin-resolver/src/finder.rs index 95d7c739e..a9c695595 100644 --- a/crates/puffin-resolver/src/finder.rs +++ b/crates/puffin-resolver/src/finder.rs @@ -3,15 +3,15 @@ //! This is similar to running `pip install` with the `--no-deps` flag. use anyhow::Result; -use distribution_filename::DistFilename; use futures::{stream, Stream, StreamExt, TryStreamExt}; use rustc_hash::FxHashMap; -use distribution_types::{Dist, File, Resolution}; +use distribution_filename::DistFilename; +use distribution_types::{Dist, IndexUrl, Resolution}; use pep440_rs::Version; use pep508_rs::{Requirement, VersionOrUrl}; -use platform_tags::{TagPriority, Tags}; -use puffin_client::{RegistryClient, SimpleMetadata}; +use platform_tags::Tags; +use puffin_client::{FlatIndex, RegistryClient, SimpleMetadata}; use puffin_interpreter::Interpreter; use puffin_normalize::PackageName; @@ -22,16 +22,23 @@ pub struct DistFinder<'a> { client: &'a RegistryClient, reporter: Option>, interpreter: &'a Interpreter, + flat_index: &'a FxHashMap>, } impl<'a> DistFinder<'a> { /// Initialize a new distribution finder. - pub fn new(tags: &'a Tags, client: &'a RegistryClient, interpreter: &'a Interpreter) -> Self { + pub fn new( + tags: &'a Tags, + client: &'a RegistryClient, + interpreter: &'a Interpreter, + flat_index: &'a FxHashMap>, + ) -> Self { Self { tags, client, reporter: None, interpreter, + flat_index, } } @@ -50,6 +57,7 @@ impl<'a> DistFinder<'a> { async fn resolve_requirement( &self, requirement: &Requirement, + flat_index: Option<&FlatIndex>, ) -> Result<(PackageName, Dist), ResolveError> { match requirement.version_or_url.as_ref() { None | Some(VersionOrUrl::VersionSpecifier(_)) => { @@ -57,17 +65,16 @@ impl<'a> DistFinder<'a> { let (index, metadata) = self.client.simple(&requirement.name).await?; // Pick a version that satisfies the requirement. - let Some(ParsedFile { filename, file }) = self.select(requirement, metadata) else { + let Some(dist) = self.select(requirement, metadata, &index, flat_index) else { return Err(ResolveError::NotFound(requirement.clone())); }; - let distribution = Dist::from_registry(filename, file, index); if let Some(reporter) = self.reporter.as_ref() { - reporter.on_progress(&distribution); + reporter.on_progress(&dist); } let normalized_name = requirement.name.clone(); - Ok((normalized_name, distribution)) + Ok((normalized_name, dist)) } Some(VersionOrUrl::Url(url)) => { // We have a URL; fetch the distribution directly. @@ -84,7 +91,9 @@ impl<'a> DistFinder<'a> { requirements: &'data [Requirement], ) -> impl Stream> + 'data { stream::iter(requirements) - .map(move |requirement| self.resolve_requirement(requirement)) + .map(move |requirement| { + self.resolve_requirement(requirement, self.flat_index.get(&requirement.name)) + }) .buffer_unordered(32) } @@ -105,10 +114,37 @@ impl<'a> DistFinder<'a> { } /// select a version that satisfies the requirement, preferring wheels to source distributions. - fn select(&self, requirement: &Requirement, metadata: SimpleMetadata) -> Option { - let mut best_version: Option = None; - let mut best_wheel: Option<(ParsedFile, TagPriority)> = None; - let mut best_sdist: Option = None; + fn select( + &self, + requirement: &Requirement, + metadata: SimpleMetadata, + index: &IndexUrl, + flat_index: Option<&FlatIndex>, + ) -> Option { + // Prioritize the flat index by initializing the "best" matches with its entries. + let matching_override = if let Some(flat_index) = flat_index { + match &requirement.version_or_url { + None => flat_index.iter().next(), + Some(VersionOrUrl::Url(_)) => None, + Some(VersionOrUrl::VersionSpecifier(specifiers)) => flat_index + .iter() + .find(|(version, _)| specifiers.contains(version)), + } + } else { + None + }; + let (mut best_version, mut best_wheel, mut best_sdist) = + if let Some((version, resolvable_dist)) = matching_override { + ( + Some(version.clone()), + resolvable_dist + .compatible_wheel() + .map(|(dist, tag_priority)| (dist.dist.clone(), *tag_priority)), + resolvable_dist.source().map(|dist| dist.dist.clone()), + ) + } else { + (None, None, None) + }; for (version, files) in metadata.into_iter().rev() { // If we iterated past the first-compatible version, break. @@ -147,10 +183,11 @@ impl<'a> DistFinder<'a> { .map_or(true, |(.., existing)| priority > *existing) { best_wheel = Some(( - ParsedFile { - filename: DistFilename::WheelFilename(wheel), + Dist::from_registry( + DistFilename::WheelFilename(wheel), file, - }, + index.clone(), + ), priority, )); } @@ -159,7 +196,7 @@ impl<'a> DistFinder<'a> { // Find the most-compatible sdist, if no wheel was found. if best_wheel.is_none() { - for (sdist, file) in files.source_dists { + for (source_dist, file) in files.source_dists { // Only add dists compatible with the python version. // This is relevant for source dists which give no other indication of their // compatibility and wheels which may be tagged `py3-none-any` but @@ -174,11 +211,12 @@ impl<'a> DistFinder<'a> { continue; } - best_version = Some(sdist.version.clone()); - best_sdist = Some(ParsedFile { - filename: DistFilename::SourceDistFilename(sdist), + best_version = Some(source_dist.version.clone()); + best_sdist = Some(Dist::from_registry( + DistFilename::SourceDistFilename(source_dist), file, - }); + index.clone(), + )); } } } @@ -187,14 +225,6 @@ impl<'a> DistFinder<'a> { } } -#[derive(Debug)] -struct ParsedFile { - /// The wheel or source dist filename extracted from the [`File`]. - filename: DistFilename, - /// The underlying [`File`]. - file: File, -} - pub trait Reporter: Send + Sync { /// Callback to invoke when a package is resolved to a specific distribution. fn on_progress(&self, dist: &Dist);