diff --git a/crates/puffin-cli/src/commands/pip_sync.rs b/crates/puffin-cli/src/commands/pip_sync.rs index 0eb1d86eb..a5f933077 100644 --- a/crates/puffin-cli/src/commands/pip_sync.rs +++ b/crates/puffin-cli/src/commands/pip_sync.rs @@ -139,8 +139,7 @@ pub(crate) async fn pip_sync( } else { let start = std::time::Instant::now(); - let flat_index_files = client.flat_index().await?; - let flat_index = FlatIndex::from_files(flat_index_files, tags); + let flat_index = FlatIndex::from_files(client.flat_index().await?, tags); let wheel_finder = puffin_resolver::DistFinder::new(tags, &client, venv.interpreter(), &flat_index) diff --git a/crates/puffin-client/src/flat_index.rs b/crates/puffin-client/src/flat_index.rs index a389d1fd0..6d0108a72 100644 --- a/crates/puffin-client/src/flat_index.rs +++ b/crates/puffin-client/src/flat_index.rs @@ -13,30 +13,30 @@ use pep440_rs::Version; use platform_tags::Tags; use puffin_normalize::PackageName; +pub type FlatIndexEntry = (DistFilename, File, IndexUrl); + +/// A set of [`PrioritizedDistribution`] from a `--find-links` entry, indexed by [`PackageName`] +/// and [`Version`]. #[derive(Debug, Clone, Default)] -pub struct FlatIndex(pub BTreeMap); +pub struct FlatIndex(FxHashMap); impl FlatIndex { - /// Collect all the files from `--find-links` into a override hashmap we can pass into version map creation. + /// Collect all files from a `--find-links` target into a [`FlatIndex`]. #[instrument(skip_all)] - pub fn from_files( - dists: Vec<(DistFilename, File, IndexUrl)>, - tags: &Tags, - ) -> FxHashMap { - // If we have packages of the same name from find links, gives them priority, otherwise start empty - let mut flat_index: FxHashMap = FxHashMap::default(); + pub fn from_files(dists: Vec, tags: &Tags) -> Self { + let mut flat_index = FxHashMap::default(); // Collect compatible distributions. for (filename, file, index) in dists { - let version_map = flat_index.entry(filename.name().clone()).or_default(); - Self::add_file(version_map, file, filename, tags, index); + let distributions = flat_index.entry(filename.name().clone()).or_default(); + Self::add_file(distributions, file, filename, tags, index); } - flat_index + Self(flat_index) } fn add_file( - version_map: &mut FlatIndex, + distributions: &mut FlatDistributions, file: File, filename: DistFilename, tags: &Tags, @@ -54,7 +54,7 @@ impl FlatIndex { file, index, })); - match version_map.0.entry(version) { + match distributions.0.entry(version) { Entry::Occupied(mut entry) => { entry.get_mut().insert_built(dist, None, None, priority); } @@ -71,7 +71,7 @@ impl FlatIndex { file, index, })); - match version_map.0.entry(filename.version.clone()) { + match distributions.0.entry(filename.version.clone()) { Entry::Occupied(mut entry) => { entry.get_mut().insert_source(dist, None, None); } @@ -83,7 +83,25 @@ impl FlatIndex { } } + /// Get the [`FlatDistributions`] for the given package name. + pub fn get(&self, package_name: &PackageName) -> Option<&FlatDistributions> { + self.0.get(package_name) + } +} + +/// A set of [`PrioritizedDistribution`] from a `--find-links` entry for a single package, indexed +/// by [`Version`]. +#[derive(Debug, Clone, Default)] +pub struct FlatDistributions(BTreeMap); + +impl FlatDistributions { pub fn iter(&self) -> impl Iterator { self.0.iter() } } + +impl From for BTreeMap { + fn from(distributions: FlatDistributions) -> Self { + distributions.0 + } +} diff --git a/crates/puffin-client/src/lib.rs b/crates/puffin-client/src/lib.rs index 580aa8239..83cf5e5e0 100644 --- a/crates/puffin-client/src/lib.rs +++ b/crates/puffin-client/src/lib.rs @@ -1,6 +1,6 @@ pub use cached_client::{CachedClient, CachedClientError, DataWithCachePolicy}; pub use error::Error; -pub use flat_index::FlatIndex; +pub use flat_index::{FlatDistributions, FlatIndex, FlatIndexEntry}; pub use registry_client::{ read_metadata_async, RegistryClient, RegistryClientBuilder, SimpleMetadata, VersionFiles, }; diff --git a/crates/puffin-client/src/registry_client.rs b/crates/puffin-client/src/registry_client.rs index 157680077..4d6817b6b 100644 --- a/crates/puffin-client/src/registry_client.rs +++ b/crates/puffin-client/src/registry_client.rs @@ -30,7 +30,7 @@ use pypi_types::{BaseUrl, Hashes, Metadata21, SimpleJson}; use crate::html::SimpleHtml; use crate::remote_metadata::wheel_metadata_from_remote_zip; -use crate::{CachedClient, CachedClientError, Error}; +use crate::{CachedClient, CachedClientError, Error, FlatIndexEntry}; /// A builder for an [`RegistryClient`]. #[derive(Debug, Clone)] @@ -118,7 +118,7 @@ impl RegistryClient { /// Read the directories and flat remote indexes from `--find-links`. #[allow(clippy::result_large_err)] - pub async fn flat_index(&self) -> Result, Error> { + pub async fn flat_index(&self) -> Result, Error> { let mut dists = Vec::new(); // TODO(konstin): Parallelize reads over flat indexes. for flat_index in self.index_locations.flat_indexes() { @@ -144,7 +144,7 @@ impl RegistryClient { } /// Read a flat remote index from a `--find-links` URL. - async fn read_flat_url(&self, url: &Url) -> Result, Error> { + async fn read_flat_url(&self, url: &Url) -> Result, Error> { let cache_entry = self.cache.entry( CacheBucket::FlatIndex, "html", @@ -198,9 +198,7 @@ impl RegistryClient { } /// Read a flat remote index from a `--find-links` directory. - fn read_flat_index_dir( - path: &PathBuf, - ) -> Result, io::Error> { + fn read_flat_index_dir(path: &PathBuf) -> Result, io::Error> { // Absolute paths are required for the URL conversion. let path = fs_err::canonicalize(path)?; let url = Url::from_directory_path(&path).expect("URL is already absolute"); diff --git a/crates/puffin-dev/src/install_many.rs b/crates/puffin-dev/src/install_many.rs index 738bd90c5..6e59d3358 100644 --- a/crates/puffin-dev/src/install_many.rs +++ b/crates/puffin-dev/src/install_many.rs @@ -18,7 +18,7 @@ use pep508_rs::Requirement; use platform_host::Platform; use platform_tags::Tags; use puffin_cache::{Cache, CacheArgs}; -use puffin_client::{RegistryClient, RegistryClientBuilder}; +use puffin_client::{FlatIndex, RegistryClient, RegistryClientBuilder}; use puffin_dispatch::BuildDispatch; use puffin_distribution::RegistryWheelIndex; use puffin_installer::Downloader; @@ -104,7 +104,7 @@ async fn install_chunk( index_locations: &IndexLocations, ) -> Result<()> { let resolution: Vec<_> = - DistFinder::new(tags, client, venv.interpreter(), &FxHashMap::default()) + DistFinder::new(tags, client, venv.interpreter(), &FlatIndex::default()) .resolve_stream(requirements) .collect() .await; diff --git a/crates/puffin-resolver/src/finder.rs b/crates/puffin-resolver/src/finder.rs index e1b31ad4a..7d73ab676 100644 --- a/crates/puffin-resolver/src/finder.rs +++ b/crates/puffin-resolver/src/finder.rs @@ -10,7 +10,7 @@ use distribution_filename::DistFilename; use distribution_types::{Dist, IndexUrl, Resolution}; use pep508_rs::{Requirement, VersionOrUrl}; use platform_tags::Tags; -use puffin_client::{FlatIndex, RegistryClient, SimpleMetadata}; +use puffin_client::{FlatDistributions, FlatIndex, RegistryClient, SimpleMetadata}; use puffin_interpreter::Interpreter; use puffin_normalize::PackageName; @@ -21,7 +21,7 @@ pub struct DistFinder<'a> { client: &'a RegistryClient, reporter: Option>, interpreter: &'a Interpreter, - flat_index: &'a FxHashMap, + flat_index: &'a FlatIndex, } impl<'a> DistFinder<'a> { @@ -30,7 +30,7 @@ impl<'a> DistFinder<'a> { tags: &'a Tags, client: &'a RegistryClient, interpreter: &'a Interpreter, - flat_index: &'a FxHashMap, + flat_index: &'a FlatIndex, ) -> Self { Self { tags, @@ -56,7 +56,7 @@ impl<'a> DistFinder<'a> { async fn resolve_requirement( &self, requirement: &Requirement, - flat_index: Option<&FlatIndex>, + flat_index: Option<&FlatDistributions>, ) -> Result<(PackageName, Dist), ResolveError> { match requirement.version_or_url.as_ref() { None | Some(VersionOrUrl::VersionSpecifier(_)) => { @@ -118,7 +118,7 @@ impl<'a> DistFinder<'a> { requirement: &Requirement, metadata: SimpleMetadata, index: &IndexUrl, - flat_index: Option<&FlatIndex>, + flat_index: Option<&FlatDistributions>, ) -> Option { // Prioritize the flat index by initializing the "best" matches with its entries. let matching_override = if let Some(flat_index) = flat_index { diff --git a/crates/puffin-resolver/src/resolver/mod.rs b/crates/puffin-resolver/src/resolver/mod.rs index 65ccc58c1..74d9d3a77 100644 --- a/crates/puffin-resolver/src/resolver/mod.rs +++ b/crates/puffin-resolver/src/resolver/mod.rs @@ -24,7 +24,7 @@ use distribution_types::{ use pep440_rs::{Version, VersionSpecifiers, MIN_VERSION}; use pep508_rs::{MarkerEnvironment, Requirement}; use platform_tags::Tags; -use puffin_client::RegistryClient; +use puffin_client::{FlatIndex, RegistryClient}; use puffin_distribution::DistributionDatabase; use puffin_interpreter::Interpreter; use puffin_normalize::PackageName; @@ -89,6 +89,7 @@ impl<'a, Context: BuildContext + Send + Sync> Resolver<'a, DefaultResolverProvid let provider = DefaultResolverProvider::new( client, DistributionDatabase::new(build_context.cache(), tags, client, build_context), + FlatIndex::from_files(client.flat_index().await?, tags), tags, PythonRequirement::new(interpreter, markers), options.exclude_newer, @@ -97,8 +98,7 @@ impl<'a, Context: BuildContext + Send + Sync> Resolver<'a, DefaultResolverProvid .iter() .chain(manifest.constraints.iter()) .collect(), - ) - .await?; + ); Ok(Self::new_custom_io( manifest, options, diff --git a/crates/puffin-resolver/src/resolver/provider.rs b/crates/puffin-resolver/src/resolver/provider.rs index d38cf9791..c2b9a7fdf 100644 --- a/crates/puffin-resolver/src/resolver/provider.rs +++ b/crates/puffin-resolver/src/resolver/provider.rs @@ -3,7 +3,6 @@ use std::future::Future; use anyhow::Result; use chrono::{DateTime, Utc}; use futures::FutureExt; -use rustc_hash::FxHashMap; use url::Url; use distribution_types::Dist; @@ -46,10 +45,12 @@ pub trait ResolverProvider: Send + Sync { /// The main IO backend for the resolver, which does cached requests network requests using the /// [`RegistryClient`] and [`DistributionDatabase`]. pub struct DefaultResolverProvider<'a, Context: BuildContext + Send + Sync> { + /// The [`RegistryClient`] used to query the index. client: &'a RegistryClient, - /// These are the entries from `--find-links` that act as overrides for index responses. - flat_index: FxHashMap, + /// The [`DistributionDatabase`] used to build source distributions. fetcher: DistributionDatabase<'a, Context>, + /// These are the entries from `--find-links` that act as overrides for index responses. + flat_index: FlatIndex, tags: &'a Tags, python_requirement: PythonRequirement<'a>, exclude_newer: Option>, @@ -58,26 +59,24 @@ pub struct DefaultResolverProvider<'a, Context: BuildContext + Send + Sync> { impl<'a, Context: BuildContext + Send + Sync> DefaultResolverProvider<'a, Context> { /// Reads the flat index entries and builds the provider. - pub async fn new( + pub fn new( client: &'a RegistryClient, fetcher: DistributionDatabase<'a, Context>, + flat_index: FlatIndex, tags: &'a Tags, python_requirement: PythonRequirement<'a>, exclude_newer: Option>, allowed_yanks: AllowedYanks, - ) -> Result { - let flat_index_dists = client.flat_index().await?; - let flat_index = FlatIndex::from_files(flat_index_dists, tags); - - Ok(Self { + ) -> Self { + Self { client, - flat_index, fetcher, + flat_index, tags, python_requirement, exclude_newer, allowed_yanks, - }) + } } } @@ -88,7 +87,6 @@ impl<'a, Context: BuildContext + Send + Sync> ResolverProvider &'io self, package_name: &'io PackageName, ) -> impl Future + Send + 'io { - let flat_index_override = self.flat_index.get(package_name).cloned(); self.client .simple(package_name) .map(move |result| match result { @@ -100,10 +98,10 @@ impl<'a, Context: BuildContext + Send + Sync> ResolverProvider &self.python_requirement, &self.allowed_yanks, self.exclude_newer.as_ref(), - flat_index_override, + self.flat_index.get(package_name).cloned(), )), Err(err @ puffin_client::Error::PackageNotFound(_)) => { - if let Some(flat_index) = flat_index_override { + if let Some(flat_index) = self.flat_index.get(package_name).cloned() { Ok(VersionMap::from(flat_index)) } else { Err(err) diff --git a/crates/puffin-resolver/src/version_map.rs b/crates/puffin-resolver/src/version_map.rs index b3c66517f..6e87dbf94 100644 --- a/crates/puffin-resolver/src/version_map.rs +++ b/crates/puffin-resolver/src/version_map.rs @@ -8,7 +8,7 @@ use distribution_filename::DistFilename; use distribution_types::{Dist, IndexUrl, PrioritizedDistribution, ResolvableDist}; use pep440_rs::Version; use platform_tags::Tags; -use puffin_client::{FlatIndex, SimpleMetadata}; +use puffin_client::{FlatDistributions, SimpleMetadata}; use puffin_normalize::PackageName; use puffin_warnings::warn_user_once; use pypi_types::{Hashes, Yanked}; @@ -32,11 +32,11 @@ impl VersionMap { python_requirement: &PythonRequirement, allowed_yanks: &AllowedYanks, exclude_newer: Option<&DateTime>, - flat_index: Option, + flat_index: Option, ) -> Self { // If we have packages of the same name from find links, gives them priority, otherwise start empty let mut version_map: BTreeMap = - flat_index.map(|overrides| overrides.0).unwrap_or_default(); + flat_index.map(Into::into).unwrap_or_default(); // Collect compatible distributions. for (version, files) in metadata { @@ -155,8 +155,8 @@ impl VersionMap { } } -impl From for VersionMap { - fn from(flat_index: FlatIndex) -> Self { - Self(flat_index.0) +impl From for VersionMap { + fn from(flat_index: FlatDistributions) -> Self { + Self(flat_index.into()) } }