diff --git a/crates/puffin-installer/src/plan.rs b/crates/puffin-installer/src/plan.rs index 4f9745ba9..f93ffb2a9 100644 --- a/crates/puffin-installer/src/plan.rs +++ b/crates/puffin-installer/src/plan.rs @@ -87,10 +87,22 @@ impl InstallPlan { // Identify any locally-available distributions that satisfy the requirement. match requirement.version_or_url.as_ref() { - None | Some(VersionOrUrl::VersionSpecifier(_)) => { - if let Some(distribution) = registry_index - .get(&requirement.name) - .filter(|dist| requirement.is_satisfied_by(&dist.filename.version)) + None => { + if let Some((_version, distribution)) = + registry_index.by_name(&requirement.name).next() + { + debug!("Requirement already cached: {distribution}"); + local.push(CachedDist::Registry(distribution.clone())); + continue; + } + } + Some(VersionOrUrl::VersionSpecifier(specifier)) => { + if let Some((_version, distribution)) = registry_index + .by_name(&requirement.name) + .find(|(version, dist)| { + specifier.contains(version) + && requirement.is_satisfied_by(&dist.filename.version) + }) { debug!("Requirement already cached: {distribution}"); local.push(CachedDist::Registry(distribution.clone())); diff --git a/crates/puffin-installer/src/registry_index.rs b/crates/puffin-installer/src/registry_index.rs index e782497b8..fbe7e41c9 100644 --- a/crates/puffin-installer/src/registry_index.rs +++ b/crates/puffin-installer/src/registry_index.rs @@ -1,9 +1,10 @@ -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use fs_err as fs; use tracing::warn; use distribution_types::{CachedRegistryDist, Metadata}; +use pep440_rs::Version; use platform_tags::Tags; use puffin_cache::{Cache, CacheBucket, WheelCache}; use puffin_normalize::PackageName; @@ -11,12 +12,12 @@ use pypi_types::IndexUrls; /// A local index of distributions that originate from a registry, like `PyPI`. #[derive(Debug, Default)] -pub struct RegistryIndex(HashMap); +pub struct RegistryIndex(HashMap>); impl RegistryIndex { /// Build an index of cached distributions from a directory. pub fn try_from_directory(cache: &Cache, tags: &Tags, index_urls: &IndexUrls) -> Self { - let mut index: HashMap = HashMap::new(); + let mut index: HashMap> = HashMap::new(); for index_url in index_urls { let wheel_dir = cache @@ -28,34 +29,36 @@ impl RegistryIndex { }; for entry in dir { - let (path, file_type) = - match entry.and_then(|entry| Ok((entry.path(), entry.file_type()?))) { - Ok((path, file_type)) => (path, file_type), - Err(err) => { - warn!( - "Failed to read entry of cache at {}: {}", - cache.root().display(), - err - ); - continue; - } - }; - if !file_type.is_dir() { - continue; - } + let path = match entry.map(|entry| entry.path()) { + Ok(path) => path, + Err(err) => { + warn!( + "Failed to read entry of cache at {}: {}", + cache.root().display(), + err + ); + continue; + } + }; match CachedRegistryDist::try_from_path(&path) { Ok(None) => {} Ok(Some(dist_info)) => { // Pick the wheel with the highest priority let compatibility = dist_info.filename.compatibility(tags); - if let Some(existing) = index.get_mut(dist_info.name()) { + if let Some(existing) = index + .get_mut(dist_info.name()) + .and_then(|package| package.get_mut(&dist_info.filename.version)) + { // Override if we have better compatibility if compatibility > existing.filename.compatibility(tags) { *existing = dist_info; } } else if compatibility.is_some() { - index.insert(dist_info.name().clone(), dist_info); + index + .entry(dist_info.name().clone()) + .or_default() + .insert(dist_info.filename.version.clone(), dist_info); } } Err(err) => { @@ -76,7 +79,13 @@ impl RegistryIndex { } /// Returns a distribution from the index, if it exists. - pub fn get(&self, name: &PackageName) -> Option<&CachedRegistryDist> { - self.0.get(name) + pub fn by_name( + &self, + name: &PackageName, + ) -> impl Iterator { + // Using static to extend the lifetime + static DEFAULT_MAP: BTreeMap = BTreeMap::new(); + // We should only query this + self.0.get(name).unwrap_or(&DEFAULT_MAP).iter().rev() } }