diff --git a/crates/rattler_installs_packages/src/package_database.rs b/crates/rattler_installs_packages/src/package_database.rs index dbfa2967..c923fee8 100644 --- a/crates/rattler_installs_packages/src/package_database.rs +++ b/crates/rattler_installs_packages/src/package_database.rs @@ -156,11 +156,15 @@ impl PackageDb { // We have exhausted all options to read the metadata from the cache. We'll have to hit the // network to get to the information. - // TODO: PEP 658 support - // Get the information from the first artifact. We assume the metadata is consistent across // all matching artifacts if let Some(artifact_info) = matching_artifacts.next() { + // Retrieve the metadata instead of the entire wheel + // If the dist-info is available separately, we can use that instead + if artifact_info.dist_info_metadata.available { + return self.get_pep658_metadata::(artifact_info).await; + } + let body = self .http .request( @@ -197,6 +201,31 @@ impl PackageDb { ); } + /// Retrieve the PEP658 metadata for the given artifact. + /// This assumes that the metadata is available in the repository + /// This can be checked with the ArtifactInfo + async fn get_pep658_metadata<'a, A: MetadataArtifact>( + &self, + artifact_info: &'a ArtifactInfo, + ) -> miette::Result<(&'a ArtifactInfo, A::Metadata)> { + // Turn into PEP658 compliant URL + let mut url = artifact_info.url.clone(); + url.set_path(&url.path().replace(".whl", ".whl.metadata")); + + let mut bytes = Vec::new(); + self.http + .request(url, Method::GET, HeaderMap::default(), CacheMode::Default) + .await? + .into_body() + .read_to_end(&mut bytes) + .await + .into_diagnostic()?; + + let metadata = A::parse_metadata(&bytes)?; + self.put_metadata_in_cache(artifact_info, &bytes)?; + Ok((artifact_info, metadata)) + } + /// Get all package names in the index. pub async fn get_package_names(&self) -> miette::Result> { let index_url = self.index_urls.first(); @@ -341,6 +370,36 @@ mod test { .await .unwrap(); } + + #[tokio::test] + async fn test_pep658() { + let cache_dir = TempDir::new().unwrap(); + let package_db = PackageDb::new( + Client::new(), + &[Url::parse("https://pypi.org/simple/").unwrap()], + cache_dir.path().into(), + ) + .unwrap(); + + // Get all the artifacts + let artifacts = package_db + .available_artifacts("numpy".parse::().unwrap()) + .await + .unwrap(); + + // Get the artifact with dist-info attribute + let artifact_info = artifacts + .iter() + .flat_map(|(_, artifacts)| artifacts.iter()) + // This signifies that a PEP658 metadata file is available + .find(|a| a.dist_info_metadata.available) + .unwrap(); + + let (_artifact, _metadata) = package_db + .get_pep658_metadata::(&artifact_info) + .await + .unwrap(); + } } #[derive(Debug, Diagnostic)]