From b64e28bb14b57a91594e0d128b135077c48c07d6 Mon Sep 17 00:00:00 2001 From: konstin Date: Thu, 3 Jul 2025 22:32:45 +0200 Subject: [PATCH] Add variants json to version map lazy --- crates/uv-client/src/registry_client.rs | 36 ++++++++++++++++--- crates/uv-distribution-types/src/lib.rs | 8 +++++ .../src/prioritized_distribution.rs | 12 ++++--- .../uv-distribution-types/src/variant_json.rs | 11 +++++- crates/uv-resolver/src/flat_index.rs | 18 ++++++---- crates/uv-resolver/src/version_map.rs | 20 ++++++++--- crates/uv/src/commands/pip/latest.rs | 2 +- 7 files changed, 86 insertions(+), 21 deletions(-) diff --git a/crates/uv-client/src/registry_client.rs b/crates/uv-client/src/registry_client.rs index b1d62edb1..ea6ff8da6 100644 --- a/crates/uv-client/src/registry_client.rs +++ b/crates/uv-client/src/registry_client.rs @@ -1058,18 +1058,21 @@ impl VersionFiles { fn push(&mut self, filename: IndexEntryFilename, file: File) { match filename { IndexEntryFilename::DistFilename(DistFilename::WheelFilename(name)) => { - self.wheels.push(VersionWheel { name, file }) + self.wheels.push(VersionWheel { name, file }); } IndexEntryFilename::DistFilename(DistFilename::SourceDistFilename(name)) => { self.source_dists.push(VersionSourceDist { name, file }); } IndexEntryFilename::VariantJson(variants_json) => { - self.variant_jsons.push(VersionVariantJson { name: variants_json, file }); + self.variant_jsons.push(VersionVariantJson { + name: variants_json, + file, + }); } } } - pub fn all(self) -> impl Iterator { + pub fn dists(self) -> impl Iterator { self.source_dists .into_iter() .map(|VersionSourceDist { name, file }| (DistFilename::SourceDistFilename(name), file)) @@ -1079,6 +1082,30 @@ impl VersionFiles { .map(|VersionWheel { name, file }| (DistFilename::WheelFilename(name), file)), ) } + + pub fn all(self) -> impl Iterator { + self.source_dists + .into_iter() + .map(|VersionSourceDist { name, file }| { + ( + IndexEntryFilename::DistFilename(DistFilename::SourceDistFilename(name)), + file, + ) + }) + .chain(self.wheels.into_iter().map(|VersionWheel { name, file }| { + ( + IndexEntryFilename::DistFilename(DistFilename::WheelFilename(name)), + file, + ) + })) + .chain( + self.variant_jsons + .into_iter() + .map(|VersionVariantJson { name, file }| { + (IndexEntryFilename::VariantJson(name), file) + }), + ) + } } #[derive(Debug, rkyv::Archive, rkyv::Deserialize, rkyv::Serialize)] @@ -1126,7 +1153,8 @@ impl SimpleMetadata { // Group the distributions by version and kind for file in files { - let Some(filename) = IndexEntryFilename::try_from_filename(&file.filename, package_name) + let Some(filename) = + IndexEntryFilename::try_from_filename(&file.filename, package_name) else { warn!("Skipping file for {package_name}: {}", file.filename); continue; diff --git a/crates/uv-distribution-types/src/lib.rs b/crates/uv-distribution-types/src/lib.rs index 66f67d39a..b6da22239 100644 --- a/crates/uv-distribution-types/src/lib.rs +++ b/crates/uv-distribution-types/src/lib.rs @@ -1465,6 +1465,14 @@ impl Identifier for BuildableSource<'_> { } } +/// A built distribution (wheel) that exists in a registry, like `PyPI`. +#[derive(Debug, Clone, Hash, PartialEq, Eq)] +pub struct RegistryVariantsJson { + pub filename: VariantsJson, + pub file: Box, + pub index: IndexUrl, +} + #[cfg(test)] mod test { use crate::{BuiltDist, Dist, RemoteSource, SourceDist, UrlString}; diff --git a/crates/uv-distribution-types/src/prioritized_distribution.rs b/crates/uv-distribution-types/src/prioritized_distribution.rs index d68568eea..279854f26 100644 --- a/crates/uv-distribution-types/src/prioritized_distribution.rs +++ b/crates/uv-distribution-types/src/prioritized_distribution.rs @@ -13,7 +13,7 @@ use uv_variants::VariantPriority; use crate::{ File, InstalledDist, KnownPlatform, RegistryBuiltDist, RegistryBuiltWheel, RegistrySourceDist, - ResolvedDistRef, VariantsJson, + RegistryVariantsJson, ResolvedDistRef, }; /// A collection of distributions that have been filtered by relevance. @@ -31,7 +31,7 @@ struct PrioritizedDistInner { /// The set of all wheels associated with this distribution. wheels: Vec<(RegistryBuiltWheel, WheelCompatibility)>, /// The `variants.json` file associated with the package version. - variants_json: Option, + variants_json: Option, /// The hashes for each distribution. hashes: Vec, /// The set of supported platforms for the distribution, described in terms of their markers. @@ -376,7 +376,7 @@ impl PrioritizedDist { } /// Create a new [`PrioritizedDist`] from the `variants.json`. - pub fn from_variant_json(variant_json: VariantsJson) -> Self { + pub fn from_variant_json(variant_json: RegistryVariantsJson) -> Self { Self(Box::new(PrioritizedDistInner { markers: MarkerTree::TRUE, best_wheel_index: None, @@ -440,7 +440,7 @@ impl PrioritizedDist { } } - pub fn insert_variant_json(&mut self, variant_json: VariantsJson) { + pub fn insert_variant_json(&mut self, variant_json: RegistryVariantsJson) { debug_assert!( self.0.variants_json.is_none(), "The variants.json filename is unique" @@ -448,6 +448,10 @@ impl PrioritizedDist { self.0.variants_json = Some(variant_json); } + pub fn variants_json(&self) -> Option<&RegistryVariantsJson> { + self.0.variants_json.as_ref() + } + /// Return the highest-priority distribution for the package version, if any. pub fn get(&self) -> Option { let best_wheel = self.0.best_wheel_index.map(|i| &self.0.wheels[i]); diff --git a/crates/uv-distribution-types/src/variant_json.rs b/crates/uv-distribution-types/src/variant_json.rs index 447d1a29a..b96b87af9 100644 --- a/crates/uv-distribution-types/src/variant_json.rs +++ b/crates/uv-distribution-types/src/variant_json.rs @@ -15,7 +15,16 @@ pub enum VariantsJsonError { /// A `--variants.json` file. #[derive( - Debug, Clone, PartialEq, Eq, PartialOrd, Ord, rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, + Debug, + Clone, + Hash, + PartialEq, + Eq, + PartialOrd, + Ord, + rkyv::Archive, + rkyv::Deserialize, + rkyv::Serialize, )] #[rkyv(derive(Debug))] pub struct VariantsJson { diff --git a/crates/uv-resolver/src/flat_index.rs b/crates/uv-resolver/src/flat_index.rs index 0de5f189d..c4b14966c 100644 --- a/crates/uv-resolver/src/flat_index.rs +++ b/crates/uv-resolver/src/flat_index.rs @@ -9,8 +9,8 @@ use uv_configuration::BuildOptions; use uv_distribution_filename::{DistFilename, SourceDistFilename, WheelFilename}; use uv_distribution_types::{ File, HashComparison, HashPolicy, IncompatibleSource, IncompatibleWheel, IndexEntryFilename, - IndexUrl, PrioritizedDist, RegistryBuiltWheel, RegistrySourceDist, SourceDistCompatibility, - WheelCompatibility, + IndexUrl, PrioritizedDist, RegistryBuiltWheel, RegistrySourceDist, RegistryVariantsJson, + SourceDistCompatibility, WheelCompatibility, }; use uv_normalize::PackageName; use uv_pep440::Version; @@ -176,13 +176,19 @@ impl FlatDistributions { } } } - IndexEntryFilename::VariantJson(variant_json) => { - match self.0.entry(variant_json.version.clone()) { + IndexEntryFilename::VariantJson(variants_json) => { + let version = variants_json.version.clone(); + let registry_variants_json = RegistryVariantsJson { + filename: variants_json, + file: Box::new(file), + index, + }; + match self.0.entry(version) { Entry::Occupied(mut entry) => { - entry.get_mut().insert_variant_json(variant_json); + entry.get_mut().insert_variant_json(registry_variants_json); } Entry::Vacant(entry) => { - entry.insert(PrioritizedDist::from_variant_json(variant_json)); + entry.insert(PrioritizedDist::from_variant_json(registry_variants_json)); } } } diff --git a/crates/uv-resolver/src/version_map.rs b/crates/uv-resolver/src/version_map.rs index 25dfb4887..b19229a41 100644 --- a/crates/uv-resolver/src/version_map.rs +++ b/crates/uv-resolver/src/version_map.rs @@ -10,9 +10,9 @@ use uv_client::{FlatIndexEntry, OwnedArchive, SimpleMetadata, VersionFiles}; use uv_configuration::BuildOptions; use uv_distribution_filename::{DistFilename, WheelFilename}; use uv_distribution_types::{ - HashComparison, IncompatibleSource, IncompatibleWheel, IndexUrl, PrioritizedDist, - RegistryBuiltWheel, RegistrySourceDist, RequiresPython, SourceDistCompatibility, - WheelCompatibility, + HashComparison, IncompatibleSource, IncompatibleWheel, IndexEntryFilename, IndexUrl, + PrioritizedDist, RegistryBuiltWheel, RegistrySourceDist, RegistryVariantsJson, RequiresPython, + SourceDistCompatibility, WheelCompatibility, }; use uv_normalize::PackageName; use uv_pep440::Version; @@ -448,7 +448,7 @@ impl VersionMapLazy { let yanked = file.yanked.as_deref(); let hashes = file.hashes.clone(); match filename { - DistFilename::WheelFilename(filename) => { + IndexEntryFilename::DistFilename(DistFilename::WheelFilename(filename)) => { let compatibility = self.wheel_compatibility( &filename, &filename.name, @@ -465,7 +465,9 @@ impl VersionMapLazy { }; priority_dist.insert_built(dist, hashes, compatibility); } - DistFilename::SourceDistFilename(filename) => { + IndexEntryFilename::DistFilename(DistFilename::SourceDistFilename( + filename, + )) => { let compatibility = self.source_dist_compatibility( &filename.name, &filename.version, @@ -484,6 +486,14 @@ impl VersionMapLazy { }; priority_dist.insert_source(dist, hashes, compatibility); } + IndexEntryFilename::VariantJson(filename) => { + let variant_json = RegistryVariantsJson { + filename, + file: Box::new(file), + index: self.index.clone(), + }; + priority_dist.insert_variant_json(variant_json); + } } } if priority_dist.is_empty() { diff --git a/crates/uv/src/commands/pip/latest.rs b/crates/uv/src/commands/pip/latest.rs index 25da8466c..8018465f4 100644 --- a/crates/uv/src/commands/pip/latest.rs +++ b/crates/uv/src/commands/pip/latest.rs @@ -68,7 +68,7 @@ impl LatestClient<'_> { // Determine whether there's a compatible wheel and/or source distribution. let mut best = None; - for (filename, file) in files.all() { + for (filename, file) in files.dists() { // Skip distributions uploaded after the cutoff. if let Some(exclude_newer) = self.exclude_newer { match file.upload_time_utc_ms.as_ref() {