Feat/pep658

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::<A>(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<Vec<String>> {
         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::<PackageName>().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::<Wheel>(&artifact_info)
+            .await
+            .unwrap();
+    }
 }
 
 #[derive(Debug, Diagnostic)]