De-duplicate result handling in Simple API responses (#10449)

## Summary

See: https://github.com/astral-sh/uv/pull/10432#issuecomment-2581084234
This commit is contained in:
Charlie Marsh 2025-01-09 17:45:13 -05:00 committed by GitHub
parent 22222e945f
commit c5e536f0ec
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -253,80 +253,36 @@ impl RegistryClient {
// If we're searching for the first index that contains the package, fetch serially.
IndexStrategy::FirstIndex => {
for index in it {
match self.simple_single_index(package_name, index).await {
Ok(metadata) => {
results.push((index, metadata));
break;
}
Err(err) => match err.into_kind() {
// The package could not be found in the remote index.
ErrorKind::WrappedReqwestError(url, err) => match err.status() {
Some(StatusCode::NOT_FOUND) => {}
Some(StatusCode::UNAUTHORIZED) => {
capabilities.set_unauthorized(index.clone());
}
Some(StatusCode::FORBIDDEN) => {
capabilities.set_forbidden(index.clone());
}
_ => return Err(ErrorKind::WrappedReqwestError(url, err).into()),
},
// The package is unavailable due to a lack of connectivity.
ErrorKind::Offline(_) => {}
// The package could not be found in the local index.
ErrorKind::FileNotFound(_) => {}
err => return Err(err.into()),
},
};
if let Some(metadata) = self
.simple_single_index(package_name, index, capabilities)
.await?
{
results.push((index, metadata));
break;
}
}
}
// Otherwise, fetch concurrently.
IndexStrategy::UnsafeBestMatch | IndexStrategy::UnsafeFirstMatch => {
let fetches = futures::stream::iter(it)
// TODO(charlie): Respect concurrency limits.
results = futures::stream::iter(it)
.map(|index| async move {
match self.simple_single_index(package_name, index).await {
Ok(metadata) => Ok(Some((index, metadata))),
Err(err) => match err.into_kind() {
// The package could not be found in the remote index.
ErrorKind::WrappedReqwestError(url, err) => match err.status() {
Some(StatusCode::NOT_FOUND) => Ok(None),
Some(StatusCode::UNAUTHORIZED) => {
capabilities.set_unauthorized(index.clone());
Ok(None)
}
Some(StatusCode::FORBIDDEN) => {
capabilities.set_forbidden(index.clone());
Ok(None)
}
_ => Err(ErrorKind::WrappedReqwestError(url, err).into()),
},
// The package is unavailable due to a lack of connectivity.
ErrorKind::Offline(_) => Ok(None),
// The package could not be found in the local index.
ErrorKind::FileNotFound(_) => Ok(None),
err => Err(err.into()),
},
let metadata = self
.simple_single_index(package_name, index, capabilities)
.await?;
Ok((index, metadata))
})
.buffered(8)
.filter_map(|result: Result<_, Error>| async move {
match result {
Ok((index, Some(metadata))) => Some(Ok((index, metadata))),
Ok((_, None)) => None,
Err(err) => Some(Err(err)),
}
})
.buffered(8);
futures::pin_mut!(fetches);
while let Some(result) = fetches.next().await {
match result {
Ok(Some((index, metadata))) => {
results.push((index, metadata));
}
Ok(None) => continue,
Err(err) => return Err(err),
}
}
.try_collect::<Vec<_>>()
.await?;
}
}
@ -346,11 +302,14 @@ impl RegistryClient {
///
/// The index can either be a PEP 503-compatible remote repository, or a local directory laid
/// out in the same format.
///
/// Returns `Ok(None)` if the package is not found in the index.
async fn simple_single_index(
&self,
package_name: &PackageName,
index: &IndexUrl,
) -> Result<OwnedArchive<SimpleMetadata>, Error> {
capabilities: &IndexCapabilities,
) -> Result<Option<OwnedArchive<SimpleMetadata>>, Error> {
// Format the URL for PyPI.
let mut url: Url = index.clone().into();
url.path_segments_mut()
@ -377,11 +336,38 @@ impl RegistryClient {
Connectivity::Offline => CacheControl::AllowStale,
};
if matches!(index, IndexUrl::Path(_)) {
let result = if matches!(index, IndexUrl::Path(_)) {
self.fetch_local_index(package_name, &url).await
} else {
self.fetch_remote_index(package_name, &url, &cache_entry, cache_control)
.await
};
match result {
Ok(metadata) => Ok(Some(metadata)),
Err(err) => match err.into_kind() {
// The package could not be found in the remote index.
ErrorKind::WrappedReqwestError(url, err) => match err.status() {
Some(StatusCode::NOT_FOUND) => Ok(None),
Some(StatusCode::UNAUTHORIZED) => {
capabilities.set_unauthorized(index.clone());
Ok(None)
}
Some(StatusCode::FORBIDDEN) => {
capabilities.set_forbidden(index.clone());
Ok(None)
}
_ => Err(ErrorKind::WrappedReqwestError(url, err).into()),
},
// The package is unavailable due to a lack of connectivity.
ErrorKind::Offline(_) => Ok(None),
// The package could not be found in the local index.
ErrorKind::FileNotFound(_) => Ok(None),
err => Err(err.into()),
},
}
}