Skip to content

Commit

Permalink
Allow multiple indexes
Browse files Browse the repository at this point in the history
  • Loading branch information
@charliermarsh
charliermarsh committed Apr 3, 2024
1 parent 8eaaf65 commit 8a8adef
Show file tree
Hide file tree
Showing 19 changed files with 355 additions and 105 deletions.
1 change: 1 addition & 0 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/uv-client/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ uv-auth = { workspace = true }
uv-cache = { workspace = true }
uv-fs = { workspace = true, features = ["tokio"] }
uv-normalize = { workspace = true }
uv-types = { workspace = true }
uv-version = { workspace = true }
uv-warnings = { workspace = true }
pypi-types = { workspace = true }
Expand Down
46 changes: 34 additions & 12 deletions crates/uv-client/src/registry_client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ use pypi_types::{Metadata23, SimpleJson};
use uv_auth::KeyringProvider;
use uv_cache::{Cache, CacheBucket, WheelCache};
use uv_normalize::PackageName;
use uv_types::IndexStrategy;

use crate::base_client::{BaseClient, BaseClientBuilder};
use crate::cached_client::CacheControl;
Expand All @@ -35,6 +36,7 @@ use crate::{CachedClient, CachedClientError, Error, ErrorKind};
#[derive(Debug, Clone)]
pub struct RegistryClientBuilder<'a> {
index_urls: IndexUrls,
index_strategy: IndexStrategy,
keyring_provider: KeyringProvider,
native_tls: bool,
retries: u32,
Expand All @@ -49,6 +51,7 @@ impl RegistryClientBuilder<'_> {
pub fn new(cache: Cache) -> Self {
Self {
index_urls: IndexUrls::default(),
index_strategy: IndexStrategy::default(),
keyring_provider: KeyringProvider::default(),
native_tls: false,
cache,
Expand All @@ -68,6 +71,12 @@ impl<'a> RegistryClientBuilder<'a> {
self
}

#[must_use]
pub fn index_strategy(mut self, index_strategy: IndexStrategy) -> Self {
self.index_strategy = index_strategy;
self
}

#[must_use]
pub fn keyring_provider(mut self, keyring_provider: KeyringProvider) -> Self {
self.keyring_provider = keyring_provider;
Expand Down Expand Up @@ -147,6 +156,7 @@ impl<'a> RegistryClientBuilder<'a> {

RegistryClient {
index_urls: self.index_urls,
index_strategy: self.index_strategy,
cache: self.cache,
connectivity,
client,
Expand All @@ -160,6 +170,8 @@ impl<'a> RegistryClientBuilder<'a> {
pub struct RegistryClient {
/// The index URLs to use for fetching packages.
index_urls: IndexUrls,
/// The strategy to use when fetching across multiple indexes.
index_strategy: IndexStrategy,
/// The underlying HTTP client.
client: CachedClient,
/// Used for the remote wheel METADATA cache.
Expand Down Expand Up @@ -206,17 +218,23 @@ impl RegistryClient {
pub async fn simple(
&self,
package_name: &PackageName,
) -> Result<(IndexUrl, OwnedArchive<SimpleMetadata>), Error> {
) -> Result<Vec<(IndexUrl, OwnedArchive<SimpleMetadata>)>, 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::new();
for index in it {
let result = self.simple_single_index(package_name, index).await?;
match self.simple_single_index(package_name, index).await? {
Ok(metadata) => {
results.push((index.clone(), metadata));

return match result {
Ok(metadata) => Ok((index.clone(), metadata)),
// If we're only using the first match, we can stop here.
if self.index_strategy == IndexStrategy::FirstMatch {
break;
}
}
Err(CachedClientError::Client(err)) => match err.into_kind() {
ErrorKind::Offline(_) => continue,
ErrorKind::ReqwestError(err) => {
Expand All @@ -225,20 +243,24 @@ impl RegistryClient {
{
continue;
}
Err(ErrorKind::from(err).into())
return Err(ErrorKind::from(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 {
Connectivity::Online => {
Err(ErrorKind::PackageNotFound(package_name.to_string()).into())
}
Connectivity::Offline => Err(ErrorKind::Offline(package_name.to_string()).into()),
if results.is_empty() {
return match self.connectivity {
Connectivity::Online => {
Err(ErrorKind::PackageNotFound(package_name.to_string()).into())
}
Connectivity::Offline => Err(ErrorKind::Offline(package_name.to_string()).into()),
};
}

Ok(results)
}

async fn simple_single_index(
Expand Down
15 changes: 11 additions & 4 deletions crates/uv-dev/src/resolve_many.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,10 +47,17 @@ async fn find_latest_version(
client: &RegistryClient,
package_name: &PackageName,
) -> Option<Version> {
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)
client
.simple(package_name)
.await
.ok()
.into_iter()
.flatten()
.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<()> {
Expand Down
55 changes: 32 additions & 23 deletions crates/uv-resolver/src/candidate_selector.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ impl CandidateSelector {
&'a self,
package_name: &'a PackageName,
range: &'a Range<Version>,
version_map: &'a VersionMap,
version_maps: &'a [VersionMap],
preferences: &'a Preferences,
installed_packages: &'a InstalledPackages,
exclusions: &'a Exclusions,
Expand Down Expand Up @@ -107,7 +107,10 @@ impl CandidateSelector {
}

// Check for a remote distribution that matches the preferred version
if let Some(file) = version_map.get(version) {
if let Some(file) = version_maps
.iter()
.find_map(|version_map| version_map.get(version))
{
return Some(Candidate::new(package_name, version, file));
}
}
Expand Down Expand Up @@ -163,33 +166,39 @@ impl CandidateSelector {
"selecting candidate for package {:?} with range {:?} with {} remote versions",
package_name,
range,
version_map.len()
version_maps.iter().map(VersionMap::len).sum::<usize>(),
);
match &self.resolution_strategy {
ResolutionStrategy::Highest => Self::select_candidate(
version_map.iter().rev(),
package_name,
range,
allow_prerelease,
),
ResolutionStrategy::Lowest => {
ResolutionStrategy::Highest => version_maps.iter().find_map(|version_map| {
Self::select_candidate(
version_map.iter().rev(),
package_name,
range,
allow_prerelease,
)
}),
ResolutionStrategy::Lowest => version_maps.iter().find_map(|version_map| {
Self::select_candidate(version_map.iter(), package_name, range, allow_prerelease)
}
}),
ResolutionStrategy::LowestDirect(direct_dependencies) => {
if direct_dependencies.contains(package_name) {
Self::select_candidate(
version_map.iter(),
package_name,
range,
allow_prerelease,
)
version_maps.iter().find_map(|version_map| {
Self::select_candidate(
version_map.iter(),
package_name,
range,
allow_prerelease,
)
})
} else {
Self::select_candidate(
version_map.iter().rev(),
package_name,
range,
allow_prerelease,
)
version_maps.iter().find_map(|version_map| {
Self::select_candidate(
version_map.iter().rev(),
package_name,
range,
allow_prerelease,
)
})
}
}
}
Expand Down
18 changes: 10 additions & 8 deletions crates/uv-resolver/src/error.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -209,14 +209,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(),
);
}
}
}
}
Expand Down
28 changes: 16 additions & 12 deletions crates/uv-resolver/src/resolution.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,12 +99,14 @@ impl ResolutionGraph {
if let Some(hash) = preferences.match_hashes(package_name, version) {
hashes.insert(package_name.clone(), hash.to_vec());
} else 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 hash = version_map.hashes(version);
hash.sort_unstable();
hash
});
if let VersionsResponse::Found(ref version_maps) = *versions_response {
for version_map in version_maps {
if let Some(mut hash) = version_map.hashes(version) {
hash.sort_unstable();
hashes.insert(package_name.clone(), hash);
break;
}
}
}
}

Expand All @@ -127,12 +129,14 @@ impl ResolutionGraph {
if let Some(hash) = preferences.match_hashes(package_name, version) {
hashes.insert(package_name.clone(), hash.to_vec());
} else 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 hash = version_map.hashes(version);
hash.sort_unstable();
hash
});
if let VersionsResponse::Found(ref version_maps) = *versions_response {
for version_map in version_maps {
if let Some(mut hash) = version_map.hashes(version) {
hash.sort_unstable();
hashes.insert(package_name.clone(), hash);
break;
}
}
}
}

Expand Down
17 changes: 7 additions & 10 deletions crates/uv-resolver/src/resolver/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ use std::ops::Deref;
use std::sync::Arc;

use anyhow::Result;

use dashmap::{DashMap, DashSet};
use futures::{FutureExt, StreamExt};
use itertools::Itertools;
Expand Down Expand Up @@ -34,7 +33,6 @@ use uv_normalize::PackageName;
use uv_types::{BuildContext, Constraints, InstalledPackagesProvider, Overrides};

use crate::candidate_selector::{CandidateDist, CandidateSelector};

use crate::editables::Editables;
use crate::error::ResolveError;
use crate::manifest::Manifest;
Expand All @@ -54,7 +52,7 @@ pub use crate::resolver::provider::{
use crate::resolver::reporter::Facade;
pub use crate::resolver::reporter::{BuildId, Reporter};
use crate::yanks::AllowedYanks;
use crate::{DependencyMode, Exclusions, Options, VersionMap};
use crate::{DependencyMode, Exclusions, Options};

mod index;
mod locals;
Expand Down Expand Up @@ -632,23 +630,22 @@ impl<
.ok_or(ResolveError::Unregistered)?;
self.visited.insert(package_name.clone());

let empty_version_map = VersionMap::default();
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.as_slice(),
VersionsResponse::NoIndex => {
self.unavailable_packages
.insert(package_name.clone(), UnavailablePackage::NoIndex);
&empty_version_map
&[]
}
VersionsResponse::Offline => {
self.unavailable_packages
.insert(package_name.clone(), UnavailablePackage::Offline);
&empty_version_map
&[]
}
VersionsResponse::NotFound => {
self.unavailable_packages
.insert(package_name.clone(), UnavailablePackage::NotFound);
&empty_version_map
&[]
}
};

Expand All @@ -664,7 +661,7 @@ impl<
let Some(candidate) = self.selector.select(
package_name,
range,
version_map,
version_maps,
&self.preferences,
self.installed_packages,
&self.exclusions,
Expand Down
Loading

0 comments on commit 8a8adef

Please sign in to comment.