From c5e536f0ec66515e20dacf237817914f3b56fc79 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Thu, 9 Jan 2025 17:45:13 -0500 Subject: [PATCH] De-duplicate result handling in Simple API responses (#10449) ## Summary See: https://github.com/astral-sh/uv/pull/10432#issuecomment-2581084234 --- crates/uv-client/src/registry_client.rs | 122 +++++++++++------------- 1 file changed, 54 insertions(+), 68 deletions(-) diff --git a/crates/uv-client/src/registry_client.rs b/crates/uv-client/src/registry_client.rs index d9812253a..6901123ab 100644 --- a/crates/uv-client/src/registry_client.rs +++ b/crates/uv-client/src/registry_client.rs @@ -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::>() + .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, Error> { + capabilities: &IndexCapabilities, + ) -> Result>, 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()), + }, } }