diff --git a/crates/uv-client/src/registry_client.rs b/crates/uv-client/src/registry_client.rs index be511b8e8ef4..a176f96220f7 100644 --- a/crates/uv-client/src/registry_client.rs +++ b/crates/uv-client/src/registry_client.rs @@ -183,17 +183,20 @@ impl RegistryClient { pub async fn simple( &self, package_name: &PackageName, - ) -> Result<(IndexUrl, OwnedArchive), Error> { + ) -> Result)>, Error> { let mut it = self.index_urls.indexes().peekable(); if it.peek().is_none() { return Err(ErrorKind::NoIndex(package_name.as_ref().to_string()).into()); } + let mut results = vec![]; for index in it { let result = self.simple_single_index(package_name, index).await?; - return match result { - Ok(metadata) => Ok((index.clone(), metadata)), + match result { + Ok(metadata) => { + results.push((index.clone(), metadata)); + } Err(CachedClientError::Client(err)) => match err.into_kind() { ErrorKind::Offline(_) => continue, ErrorKind::RequestError(err) => { @@ -201,13 +204,18 @@ impl RegistryClient { || err.status() == Some(StatusCode::FORBIDDEN) { continue; + } else { + return Err(ErrorKind::RequestError(err).into()); } - Err(ErrorKind::RequestError(err).into()) } - other => Err(other.into()), + other => { + return Err(other.into()); + } }, - Err(CachedClientError::Callback(err)) => Err(err), - }; + Err(CachedClientError::Callback(err)) => { + return Err(err); + } + } } match self.connectivity { diff --git a/crates/uv-dev/src/resolve_many.rs b/crates/uv-dev/src/resolve_many.rs index ebcb86cf4e9e..d13b806322f5 100644 --- a/crates/uv-dev/src/resolve_many.rs +++ b/crates/uv-dev/src/resolve_many.rs @@ -48,10 +48,15 @@ async fn find_latest_version( client: &RegistryClient, package_name: &PackageName, ) -> Option { - let (_, raw_simple_metadata) = client.simple(package_name).await.ok()?; - let simple_metadata = OwnedArchive::deserialize(&raw_simple_metadata); - let version = simple_metadata.into_iter().next()?.version; - Some(version) + let results = client.simple(package_name).await.ok()?; + + results + .into_iter() + .filter_map(|(_index, raw_simple_metadata)| { + let simple_metadata = OwnedArchive::deserialize(&raw_simple_metadata); + Some(simple_metadata.into_iter().next()?.version) + }) + .max() } pub(crate) async fn resolve_many(args: ResolveManyArgs) -> Result<()> { diff --git a/crates/uv-resolver/src/candidate_selector.rs b/crates/uv-resolver/src/candidate_selector.rs index d62db6940525..24cf943c10c0 100644 --- a/crates/uv-resolver/src/candidate_selector.rs +++ b/crates/uv-resolver/src/candidate_selector.rs @@ -89,14 +89,16 @@ impl CandidateSelector { &'a self, package_name: &'a PackageName, range: &'a Range, - version_map: &'a VersionMap, + version_maps: &'a [VersionMap], ) -> Option> { // If the package has a preference (e.g., an existing version from an existing lockfile), // and the preference satisfies the current range, use that. if let Some(version) = self.preferences.get(package_name) { if range.contains(version) { - if let Some(file) = version_map.get(version) { - return Some(Candidate::new(package_name, version, file)); + for version_map in version_maps { + if let Some(file) = version_map.get(version) { + return Some(Candidate::new(package_name, version, file)); + } } } } @@ -126,29 +128,29 @@ impl CandidateSelector { "selecting candidate for package {:?} with range {:?} with {} versions", package_name, range, - version_map.len() + version_maps + .iter() + .map(|version_map| version_map.len()) + .sum::(), ); match &self.resolution_strategy { - ResolutionStrategy::Highest => Self::select_candidate( - version_map.iter().rev(), - package_name, - range, - allow_prerelease, - ), + ResolutionStrategy::Highest => { + Self::select_highest_candidate(version_maps, package_name, range, allow_prerelease) + } ResolutionStrategy::Lowest => { - Self::select_candidate(version_map.iter(), package_name, range, allow_prerelease) + Self::select_lowest_candidate(version_maps, package_name, range, allow_prerelease) } ResolutionStrategy::LowestDirect(direct_dependencies) => { if direct_dependencies.contains(package_name) { - Self::select_candidate( - version_map.iter(), + Self::select_lowest_candidate( + version_maps, package_name, range, allow_prerelease, ) } else { - Self::select_candidate( - version_map.iter().rev(), + Self::select_highest_candidate( + version_maps, package_name, range, allow_prerelease, @@ -158,6 +160,39 @@ impl CandidateSelector { } } + fn select_highest_candidate<'a>( + version_maps: &'a [VersionMap], + package_name: &'a PackageName, + range: &Range, + allow_prerelease: AllowPreRelease, + ) -> Option> { + version_maps + .iter() + .filter_map(|version_map| { + Self::select_candidate( + version_map.iter().rev(), + package_name, + range, + allow_prerelease, + ) + }) + .max_by_key(|candidate| candidate.version) + } + + fn select_lowest_candidate<'a>( + version_maps: &'a [VersionMap], + package_name: &'a PackageName, + range: &Range, + allow_prerelease: AllowPreRelease, + ) -> Option> { + version_maps + .iter() + .filter_map(|version_map| { + Self::select_candidate(version_map.iter(), package_name, range, allow_prerelease) + }) + .min_by_key(|candidate| candidate.version) + } + /// Select the first-matching [`Candidate`] from a set of candidate versions and files, /// preferring wheels over source distributions. fn select_candidate<'a>( diff --git a/crates/uv-resolver/src/error.rs b/crates/uv-resolver/src/error.rs index 9f3857dbb28e..27a7a0ad2ad6 100644 --- a/crates/uv-resolver/src/error.rs +++ b/crates/uv-resolver/src/error.rs @@ -206,14 +206,16 @@ impl NoSolutionError { // we represent the state of the resolver at the time of failure. if visited.contains(name) { if let Some(response) = package_versions.get(name) { - if let VersionsResponse::Found(ref version_map) = *response { - available_versions.insert( - package.clone(), - version_map - .iter() - .map(|(version, _)| version.clone()) - .collect(), - ); + if let VersionsResponse::Found(ref version_maps) = *response { + for version_map in version_maps { + available_versions.insert( + package.clone(), + version_map + .iter() + .map(|(version, _)| version.clone()) + .collect(), + ); + } } } } diff --git a/crates/uv-resolver/src/finder.rs b/crates/uv-resolver/src/finder.rs index b49e36446e2b..1d4a7ea533b6 100644 --- a/crates/uv-resolver/src/finder.rs +++ b/crates/uv-resolver/src/finder.rs @@ -4,6 +4,7 @@ use anyhow::Result; use futures::{stream, Stream, StreamExt, TryStreamExt}; +use itertools::Itertools; use rustc_hash::FxHashMap; use uv_traits::NoBinary; @@ -67,11 +68,14 @@ impl<'a> DistFinder<'a> { match requirement.version_or_url.as_ref() { None | Some(VersionOrUrl::VersionSpecifier(_)) => { // Query the index(es) (cached) to get the URLs for the available files. - let (index, raw_metadata) = self.client.simple(&requirement.name).await?; - let metadata = OwnedArchive::deserialize(&raw_metadata); + let results = self.client.simple(&requirement.name).await?; + let results = results + .iter() + .map(|(index, raw_metadata)| (index, OwnedArchive::deserialize(&raw_metadata))) + .collect::>(); // Pick a version that satisfies the requirement. - let Some(dist) = self.select(requirement, metadata, &index, flat_index) else { + let Some(dist) = self.select(requirement, results, flat_index) else { return Err(ResolveError::NotFound(requirement.clone())); }; @@ -126,8 +130,7 @@ impl<'a> DistFinder<'a> { fn select( &self, requirement: &Requirement, - metadata: SimpleMetadata, - index: &IndexUrl, + metadata_tuples: Vec<(&IndexUrl, SimpleMetadata)>, flat_index: Option<&FlatDistributions>, ) -> Option { let no_binary = match self.no_binary { @@ -161,7 +164,18 @@ impl<'a> DistFinder<'a> { (None, None, None) }; - for SimpleMetadatum { version, files } in metadata.into_iter().rev() { + for (index, SimpleMetadatum { version, files }) in metadata_tuples + .into_iter() + .map(|(index, metadata)| { + metadata + .into_iter() + .map(move |metadatum| (index, metadatum)) + .rev() + }) + .kmerge_by(|(_index1, metadata1), (_index2, metadata2)| { + metadata1.version >= metadata2.version + }) + { // If we iterated past the first-compatible version, break. if best_version .as_ref() diff --git a/crates/uv-resolver/src/resolution.rs b/crates/uv-resolver/src/resolution.rs index 49ac8b61857a..ebe1d79b46ce 100644 --- a/crates/uv-resolver/src/resolution.rs +++ b/crates/uv-resolver/src/resolution.rs @@ -86,12 +86,14 @@ impl ResolutionGraph { // Add its hashes to the index. if let Some(versions_response) = packages.get(package_name) { - if let VersionsResponse::Found(ref version_map) = *versions_response { - hashes.insert(package_name.clone(), { - let mut hashes = version_map.hashes(version); - hashes.sort_unstable(); - hashes - }); + if let VersionsResponse::Found(ref version_maps) = *versions_response { + for version_map in version_maps { + hashes.insert(package_name.clone(), { + let mut hashes = version_map.hashes(version); + hashes.sort_unstable(); + hashes + }); + } } } @@ -113,12 +115,14 @@ impl ResolutionGraph { // Add its hashes to the index. if let Some(versions_response) = packages.get(package_name) { - if let VersionsResponse::Found(ref version_map) = *versions_response { - hashes.insert(package_name.clone(), { - let mut hashes = version_map.hashes(version); - hashes.sort_unstable(); - hashes - }); + if let VersionsResponse::Found(ref version_maps) = *versions_response { + for version_map in version_maps { + hashes.insert(package_name.clone(), { + let mut hashes = version_map.hashes(version); + hashes.sort_unstable(); + hashes + }); + } } } diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index c8c98a388d93..fd3824692314 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -614,8 +614,8 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> { .ok_or(ResolveError::Unregistered)?; self.visited.insert(package_name.clone()); - let version_map = match *versions_response { - VersionsResponse::Found(ref version_map) => version_map, + let version_maps = match *versions_response { + VersionsResponse::Found(ref version_maps) => version_maps, // Short-circuit if we do not find any versions for the package VersionsResponse::NoIndex => { self.unavailable_packages @@ -646,7 +646,8 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> { } // Find a version. - let Some(candidate) = self.selector.select(package_name, range, version_map) else { + let Some(candidate) = self.selector.select(package_name, range, version_maps) + else { // Short circuit: we couldn't find _any_ versions for a package. return Ok(None); }; @@ -1008,8 +1009,8 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> { .await .ok_or(ResolveError::Unregistered)?; - let version_map = match *versions_response { - VersionsResponse::Found(ref version_map) => version_map, + let version_maps = match *versions_response { + VersionsResponse::Found(ref version_maps) => version_maps, // Short-circuit if we did not find any versions for the package VersionsResponse::NoIndex => { self.unavailable_packages @@ -1033,7 +1034,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> { // Try to find a compatible version. If there aren't any compatible versions, // short-circuit and return `None`. - let Some(candidate) = self.selector.select(&package_name, &range, version_map) + let Some(candidate) = self.selector.select(&package_name, &range, version_maps) else { return Ok(None); }; diff --git a/crates/uv-resolver/src/resolver/provider.rs b/crates/uv-resolver/src/resolver/provider.rs index 929c5e1efc3d..6d579e40efc1 100644 --- a/crates/uv-resolver/src/resolver/provider.rs +++ b/crates/uv-resolver/src/resolver/provider.rs @@ -21,8 +21,8 @@ pub type WheelMetadataResult = Result<(Metadata21, Option), uv_distribution /// The response when requesting versions for a package #[derive(Debug)] pub enum VersionsResponse { - /// The package was found in the registry with the included versions - Found(VersionMap), + /// The package was found in the registry/registries with the included versions + Found(Vec), /// The package was not found in the registry NotFound, /// The package was not found in the local registry @@ -102,32 +102,39 @@ impl<'a, Context: BuildContext + Send + Sync> ResolverProvider &'io self, package_name: &'io PackageName, ) -> PackageVersionsResult { - let result = self.client.simple(package_name).await; + let results = self.client.simple(package_name).await; // If the "Simple API" request was successful, convert to `VersionMap` on the Tokio // threadpool, since it can be slow. - match result { - Ok((index, metadata)) => Ok(VersionsResponse::Found(VersionMap::from_metadata( - metadata, - package_name, - &index, - &self.tags, - &self.python_requirement, - self.exclude_newer.as_ref(), - self.flat_index.get(package_name).cloned(), - &self.no_binary, - ))), + match results { + Ok(results) => Ok(VersionsResponse::Found( + results + .into_iter() + .map(|(index, metadata)| { + VersionMap::from_metadata( + metadata, + package_name, + &index, + &self.tags, + &self.python_requirement, + self.exclude_newer.as_ref(), + self.flat_index.get(package_name).cloned(), + &self.no_binary, + ) + }) + .collect(), + )), Err(err) => match err.into_kind() { uv_client::ErrorKind::PackageNotFound(_) => { if let Some(flat_index) = self.flat_index.get(package_name).cloned() { - Ok(VersionsResponse::Found(VersionMap::from(flat_index))) + Ok(VersionsResponse::Found(vec![VersionMap::from(flat_index)])) } else { Ok(VersionsResponse::NotFound) } } uv_client::ErrorKind::NoIndex(_) => { if let Some(flat_index) = self.flat_index.get(package_name).cloned() { - Ok(VersionsResponse::Found(VersionMap::from(flat_index))) + Ok(VersionsResponse::Found(vec![VersionMap::from(flat_index)])) } else if self.flat_index.offline() { Ok(VersionsResponse::Offline) } else { @@ -136,7 +143,7 @@ impl<'a, Context: BuildContext + Send + Sync> ResolverProvider } uv_client::ErrorKind::Offline(_) => { if let Some(flat_index) = self.flat_index.get(package_name).cloned() { - Ok(VersionsResponse::Found(VersionMap::from(flat_index))) + Ok(VersionsResponse::Found(vec![VersionMap::from(flat_index)])) } else { Ok(VersionsResponse::Offline) }