Skip to content

Commit

Permalink
Feat/pep658 (#15)
Browse files Browse the repository at this point in the history
Adds support for [PEP658](https://peps.python.org/pep-0658/) metadata
files and downloads these instead of entire wheel if possible.

Added small test that tests these with numpy which definitely supports
them in the latest versions.
  • Loading branch information
tdejager authored Sep 26, 2023
1 parent 3aa5655 commit 8faecd1
Showing 1 changed file with 61 additions and 2 deletions.
63 changes: 61 additions & 2 deletions crates/rattler_installs_packages/src/package_database.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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)]
Expand Down

0 comments on commit 8faecd1

Please sign in to comment.