diff --git a/crates/uv-resolver/src/lock.rs b/crates/uv-resolver/src/lock.rs index 29f61ac87..3775be8bd 100644 --- a/crates/uv-resolver/src/lock.rs +++ b/crates/uv-resolver/src/lock.rs @@ -4,21 +4,23 @@ use std::collections::VecDeque; +use rustc_hash::FxHashMap; +use url::Url; + use distribution_filename::WheelFilename; use distribution_types::{ - BuiltDist, DirectUrlBuiltDist, DirectUrlSourceDist, DirectorySourceDist, Dist, - DistributionMetadata, FileLocation, GitSourceDist, IndexUrl, Name, PathBuiltDist, - PathSourceDist, RegistryBuiltDist, RegistrySourceDist, Resolution, ResolvedDist, ToUrlError, - VersionOrUrlRef, + BuiltDist, DirectUrlBuiltDist, DirectUrlSourceDist, DirectorySourceDist, Dist, FileLocation, + GitSourceDist, IndexUrl, PathBuiltDist, PathSourceDist, RegistryBuiltDist, RegistrySourceDist, + Resolution, ResolvedDist, ToUrlError, }; use pep440_rs::Version; use pep508_rs::{MarkerEnvironment, VerbatimUrl}; use platform_tags::{TagCompatibility, TagPriority, Tags}; use pypi_types::HashDigest; -use rustc_hash::FxHashMap; -use url::Url; use uv_normalize::PackageName; +use crate::resolution::AnnotatedDist; + #[derive(Clone, Debug, serde::Deserialize, serde::Serialize)] #[serde(into = "LockWire", try_from = "LockWire")] pub struct Lock { @@ -71,7 +73,7 @@ impl Lock { let dep_dist = self.find_by_id(&dep.id); queue.push_back(dep_dist); } - let name = PackageName::new(dist.id.name.to_string()).unwrap(); + let name = dist.id.name.clone(); let resolved_dist = ResolvedDist::Installable(dist.to_dist(marker_env, tags)); map.insert(name, resolved_dist); } @@ -202,15 +204,15 @@ pub(crate) struct Distribution { } impl Distribution { - pub(crate) fn from_resolved_dist( - resolved_dist: &ResolvedDist, + pub(crate) fn from_annotated_dist( + annotated_dist: &AnnotatedDist, ) -> Result { - let id = DistributionId::from_resolved_dist(resolved_dist); + let id = DistributionId::from_annotated_dist(annotated_dist); let mut sdist = None; let mut wheels = vec![]; - if let Some(dist) = Wheel::from_resolved_dist(resolved_dist)? { + if let Some(dist) = Wheel::from_annotated_dist(annotated_dist)? { wheels.push(dist); - } else if let Some(dist) = SourceDist::from_resolved_dist(resolved_dist)? { + } else if let Some(dist) = SourceDist::from_annotated_dist(annotated_dist)? { sdist = Some(dist); } Ok(Distribution { @@ -224,9 +226,9 @@ impl Distribution { }) } - pub(crate) fn add_dependency(&mut self, resolved_dist: &ResolvedDist) { + pub(crate) fn add_dependency(&mut self, annotated_dist: &AnnotatedDist) { self.dependencies - .push(Dependency::from_resolved_dist(resolved_dist)); + .push(Dependency::from_annotated_dist(annotated_dist)); } fn to_dist(&self, _marker_env: &MarkerEnvironment, tags: &Tags) -> Dist { @@ -294,16 +296,10 @@ pub(crate) struct DistributionId { } impl DistributionId { - fn from_resolved_dist(resolved_dist: &ResolvedDist) -> DistributionId { - let name = resolved_dist.name().clone(); - let version = match resolved_dist.version_or_url() { - VersionOrUrlRef::Version(v) => v.clone(), - // TODO: We need a way to thread the version number for these - // cases down into this routine. The version number isn't yet in a - // `ResolutionGraph`, so this will require a bit of refactoring. - VersionOrUrlRef::Url(_) => todo!(), - }; - let source = Source::from_resolved_dist(resolved_dist); + fn from_annotated_dist(annotated_dist: &AnnotatedDist) -> DistributionId { + let name = annotated_dist.metadata.name.clone(); + let version = annotated_dist.metadata.version.clone(); + let source = Source::from_resolved_dist(&annotated_dist.dist); DistributionId { name, version, @@ -596,85 +592,89 @@ pub(crate) struct SourceDist { } impl SourceDist { - fn from_resolved_dist(resolved_dist: &ResolvedDist) -> Result, LockError> { - match *resolved_dist { + fn from_annotated_dist( + annotated_dist: &AnnotatedDist, + ) -> Result, LockError> { + match annotated_dist.dist { // TODO: Do we want to try to lock already-installed distributions? // Or should we return an error? ResolvedDist::Installed(_) => todo!(), - ResolvedDist::Installable(ref dist) => SourceDist::from_dist(dist), + ResolvedDist::Installable(ref dist) => { + SourceDist::from_dist(dist, &annotated_dist.hashes) + } } } - fn from_dist(dist: &Dist) -> Result, LockError> { + fn from_dist(dist: &Dist, hashes: &[HashDigest]) -> Result, LockError> { match *dist { Dist::Built(_) => Ok(None), - Dist::Source(ref source_dist) => SourceDist::from_source_dist(source_dist).map(Some), + Dist::Source(ref source_dist) => { + SourceDist::from_source_dist(source_dist, hashes).map(Some) + } } } fn from_source_dist( source_dist: &distribution_types::SourceDist, + hashes: &[HashDigest], ) -> Result { match *source_dist { distribution_types::SourceDist::Registry(ref reg_dist) => { SourceDist::from_registry_dist(reg_dist) } distribution_types::SourceDist::DirectUrl(ref direct_dist) => { - Ok(SourceDist::from_direct_dist(direct_dist)) + Ok(SourceDist::from_direct_dist(direct_dist, hashes)) } distribution_types::SourceDist::Git(ref git_dist) => { - Ok(SourceDist::from_git_dist(git_dist)) + Ok(SourceDist::from_git_dist(git_dist, hashes)) } distribution_types::SourceDist::Path(ref path_dist) => { - Ok(SourceDist::from_path_dist(path_dist)) + Ok(SourceDist::from_path_dist(path_dist, hashes)) } distribution_types::SourceDist::Directory(ref directory_dist) => { - Ok(SourceDist::from_directory_dist(directory_dist)) + Ok(SourceDist::from_directory_dist(directory_dist, hashes)) } } } fn from_registry_dist(reg_dist: &RegistrySourceDist) -> Result { - // FIXME: Is it guaranteed that there is at least one hash? - // If not, we probably need to make this fallible. let url = reg_dist .file .url .to_url() .map_err(LockError::invalid_file_url)?; - let hash = Hash::from(reg_dist.file.hashes[0].clone()); - Ok(SourceDist { - url, - hash: Some(hash), - }) + let hash = reg_dist.file.hashes.first().cloned().map(Hash::from); + Ok(SourceDist { url, hash }) } - fn from_direct_dist(direct_dist: &DirectUrlSourceDist) -> SourceDist { + fn from_direct_dist(direct_dist: &DirectUrlSourceDist, hashes: &[HashDigest]) -> SourceDist { SourceDist { url: direct_dist.url.to_url(), - // TODO: We want a hash for the artifact at the URL. - hash: todo!(), + hash: hashes.first().cloned().map(Hash::from), } } - fn from_git_dist(git_dist: &GitSourceDist) -> SourceDist { + fn from_git_dist(git_dist: &GitSourceDist, hashes: &[HashDigest]) -> SourceDist { SourceDist { url: git_dist.url.to_url(), - hash: None, + hash: hashes.first().cloned().map(Hash::from), } } - fn from_path_dist(path_dist: &PathSourceDist) -> SourceDist { + fn from_path_dist(path_dist: &PathSourceDist, hashes: &[HashDigest]) -> SourceDist { SourceDist { url: path_dist.url.to_url(), - hash: None, + hash: hashes.first().cloned().map(Hash::from), } } - fn from_directory_dist(directory_dist: &DirectorySourceDist) -> SourceDist { + fn from_directory_dist( + directory_dist: &DirectorySourceDist, + hashes: &[HashDigest], + ) -> SourceDist { SourceDist { url: directory_dist.url.to_url(), - hash: None, + hash: hashes.first().cloned().map(Hash::from), } } } @@ -704,60 +704,59 @@ pub(crate) struct Wheel { } impl Wheel { - fn from_resolved_dist(resolved_dist: &ResolvedDist) -> Result, LockError> { - match *resolved_dist { + fn from_annotated_dist(annotated_dist: &AnnotatedDist) -> Result, LockError> { + match annotated_dist.dist { // TODO: Do we want to try to lock already-installed distributions? // Or should we return an error? ResolvedDist::Installed(_) => todo!(), - ResolvedDist::Installable(ref dist) => Wheel::from_dist(dist), + ResolvedDist::Installable(ref dist) => Wheel::from_dist(dist, &annotated_dist.hashes), } } - fn from_dist(dist: &Dist) -> Result, LockError> { + fn from_dist(dist: &Dist, hashes: &[HashDigest]) -> Result, LockError> { match *dist { - Dist::Built(ref built_dist) => Wheel::from_built_dist(built_dist).map(Some), + Dist::Built(ref built_dist) => Wheel::from_built_dist(built_dist, hashes).map(Some), Dist::Source(_) => Ok(None), } } - fn from_built_dist(built_dist: &BuiltDist) -> Result { + fn from_built_dist(built_dist: &BuiltDist, hashes: &[HashDigest]) -> Result { match *built_dist { BuiltDist::Registry(ref reg_dist) => Wheel::from_registry_dist(reg_dist), - BuiltDist::DirectUrl(ref direct_dist) => Ok(Wheel::from_direct_dist(direct_dist)), - BuiltDist::Path(ref path_dist) => Ok(Wheel::from_path_dist(path_dist)), + BuiltDist::DirectUrl(ref direct_dist) => { + Ok(Wheel::from_direct_dist(direct_dist, hashes)) + } + BuiltDist::Path(ref path_dist) => Ok(Wheel::from_path_dist(path_dist, hashes)), } } fn from_registry_dist(reg_dist: &RegistryBuiltDist) -> Result { - // FIXME: Is it guaranteed that there is at least one hash? - // If not, we probably need to make this fallible. let filename = reg_dist.filename.clone(); let url = reg_dist .file .url .to_url() .map_err(LockError::invalid_file_url)?; - let hash = Hash::from(reg_dist.file.hashes[0].clone()); + let hash = reg_dist.file.hashes.first().cloned().map(Hash::from); Ok(Wheel { url, - hash: Some(hash), + hash, filename, }) } - fn from_direct_dist(direct_dist: &DirectUrlBuiltDist) -> Wheel { + fn from_direct_dist(direct_dist: &DirectUrlBuiltDist, hashes: &[HashDigest]) -> Wheel { Wheel { url: direct_dist.url.to_url(), - // TODO: We want a hash for the artifact at the URL. - hash: todo!(), + hash: hashes.first().cloned().map(Hash::from), filename: direct_dist.filename.clone(), } } - fn from_path_dist(path_dist: &PathBuiltDist) -> Wheel { + fn from_path_dist(path_dist: &PathBuiltDist, hashes: &[HashDigest]) -> Wheel { Wheel { url: path_dist.url.to_url(), - hash: None, + hash: hashes.first().cloned().map(Hash::from), filename: path_dist.filename.clone(), } } @@ -816,8 +815,8 @@ pub(crate) struct Dependency { } impl Dependency { - fn from_resolved_dist(resolved_dist: &ResolvedDist) -> Dependency { - let id = DistributionId::from_resolved_dist(resolved_dist); + fn from_annotated_dist(annotated_dist: &AnnotatedDist) -> Dependency { + let id = DistributionId::from_annotated_dist(annotated_dist); Dependency { id } } } diff --git a/crates/uv-resolver/src/resolution.rs b/crates/uv-resolver/src/resolution.rs index 7ccaa3b97..9cba84f9b 100644 --- a/crates/uv-resolver/src/resolution.rs +++ b/crates/uv-resolver/src/resolution.rs @@ -1,5 +1,6 @@ use std::borrow::Cow; use std::collections::BTreeSet; +use std::fmt::Display; use std::hash::BuildHasherDefault; use std::rc::Rc; @@ -20,7 +21,7 @@ use distribution_types::{ use once_map::OnceMap; use pep440_rs::{Version, VersionSpecifier}; use pep508_rs::MarkerEnvironment; -use pypi_types::HashDigest; +use pypi_types::{HashDigest, Metadata23}; use uv_normalize::{ExtraName, PackageName}; use crate::dependency_provider::UvDependencyProvider; @@ -47,14 +48,40 @@ pub enum AnnotationStyle { Split, } +/// A pinned package with its resolved distribution and metadata. The [`ResolvedDist`] refers to a +/// specific distribution (e.g., a specific wheel), while the [`Metadata23`] refers to the metadata +/// for the package-version pair. +#[derive(Debug)] +pub struct AnnotatedDist { + pub dist: ResolvedDist, + pub hashes: Vec, + pub metadata: Metadata23, +} + +impl Name for AnnotatedDist { + fn name(&self) -> &PackageName { + self.dist.name() + } +} + +impl DistributionMetadata for AnnotatedDist { + fn version_or_url(&self) -> VersionOrUrlRef { + self.dist.version_or_url() + } +} + +impl Display for AnnotatedDist { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + Display::fmt(&self.dist, f) + } +} + /// A complete resolution graph in which every node represents a pinned package and every edge /// represents a dependency between two pinned packages. #[derive(Debug)] pub struct ResolutionGraph { /// The underlying graph. - petgraph: petgraph::graph::Graph, petgraph::Directed>, - /// The metadata for every distribution in this resolution. - hashes: FxHashMap>, + petgraph: petgraph::graph::Graph, petgraph::Directed>, /// The enabled extras for every distribution in this resolution. extras: FxHashMap>, /// The set of editable requirements in this resolution. @@ -186,21 +213,18 @@ impl ResolutionGraph { }; } + // Add every package to the graph. // TODO(charlie): petgraph is a really heavy and unnecessary dependency here. We should // write our own graph, given that our requirements are so simple. let mut petgraph = petgraph::graph::Graph::with_capacity(selection.len(), selection.len()); - let mut hashes = - FxHashMap::with_capacity_and_hasher(selection.len(), BuildHasherDefault::default()); - - // Add every package to the graph. let mut inverse = FxHashMap::with_capacity_and_hasher(selection.len(), BuildHasherDefault::default()); + for (package, version) in selection { match package { PubGrubPackage::Package(package_name, None, None) => { // Create the distribution. - let pinned_package = if let Some((editable, _, _)) = editables.get(package_name) - { + let dist = if let Some((editable, _, _)) = editables.get(package_name) { Dist::from_editable(package_name.clone(), editable.clone())?.into() } else { pins.get(package_name, version) @@ -208,57 +232,117 @@ impl ResolutionGraph { .clone() }; - // Add its hashes to the index, preserving those that were already present in - // the lockfile if necessary. - if let Some(digests) = preferences + // Extract the hashes, preserving those that were already present in the + // lockfile if necessary. + let hashes = if let Some(digests) = preferences .match_hashes(package_name, version) .filter(|digests| !digests.is_empty()) { - hashes.insert(package_name.clone(), digests.to_vec()); + digests.to_vec() } else if let Some(versions_response) = packages.get(package_name) { if let VersionsResponse::Found(ref version_maps) = *versions_response { - for version_map in version_maps { - if let Some(mut digests) = version_map.hashes(version) { + version_maps + .iter() + .find_map(|version_map| version_map.hashes(version)) + .map(|mut digests| { digests.sort_unstable(); - hashes.insert(package_name.clone(), digests); - break; - } - } + digests + }) + .unwrap_or_default() + } else { + vec![] } - } + } else { + vec![] + }; + + // Extract the metadata. + let metadata = if let Some((_, metadata, _)) = editables.get(package_name) { + metadata.clone() + } else { + let dist = PubGrubDistribution::from_registry(package_name, version); + + let response = distributions.get(&dist.version_id()).unwrap_or_else(|| { + panic!( + "Every package should have metadata: {:?}", + dist.version_id() + ) + }); + + let MetadataResponse::Found(archive) = &*response else { + panic!( + "Every package should have metadata: {:?}", + dist.version_id() + ) + }; + + archive.metadata.clone() + }; // Add the distribution to the graph. - let index = petgraph.add_node(pinned_package); + let index = petgraph.add_node(AnnotatedDist { + dist, + hashes, + metadata, + }); inverse.insert(package_name, index); } PubGrubPackage::Package(package_name, None, Some(url)) => { // Create the distribution. - let pinned_package = if let Some((editable, _, _)) = editables.get(package_name) - { + let dist = if let Some((editable, _, _)) = editables.get(package_name) { Dist::from_editable(package_name.clone(), editable.clone())? } else { Dist::from_url(package_name.clone(), url_to_precise(url.clone()))? }; - // Add its hashes to the index, preserving those that were already present in - // the lockfile if necessary. - if let Some(digests) = preferences + // Extract the hashes, preserving those that were already present in the + // lockfile if necessary. + let hashes = if let Some(digests) = preferences .match_hashes(package_name, version) .filter(|digests| !digests.is_empty()) { - hashes.insert(package_name.clone(), digests.to_vec()); - } else if let Some(metadata_response) = - distributions.get(&pinned_package.version_id()) - { + digests.to_vec() + } else if let Some(metadata_response) = distributions.get(&dist.version_id()) { if let MetadataResponse::Found(ref archive) = *metadata_response { let mut digests = archive.hashes.clone(); digests.sort_unstable(); - hashes.insert(package_name.clone(), digests); + digests + } else { + vec![] } - } + } else { + vec![] + }; + + // Extract the metadata. + let metadata = if let Some((_, metadata, _)) = editables.get(package_name) { + metadata.clone() + } else { + let dist = PubGrubDistribution::from_url(package_name, url); + + let response = distributions.get(&dist.version_id()).unwrap_or_else(|| { + panic!( + "Every package should have metadata: {:?}", + dist.version_id() + ) + }); + + let MetadataResponse::Found(archive) = &*response else { + panic!( + "Every package should have metadata: {:?}", + dist.version_id() + ) + }; + + archive.metadata.clone() + }; // Add the distribution to the graph. - let index = petgraph.add_node(pinned_package.into()); + let index = petgraph.add_node(AnnotatedDist { + dist: dist.into(), + hashes, + metadata, + }); inverse.insert(package_name, index); } _ => {} @@ -311,7 +395,6 @@ impl ResolutionGraph { Ok(Self { petgraph, - hashes, extras, editables, diagnostics, @@ -341,7 +424,7 @@ impl ResolutionGraph { .into_nodes_edges() .0 .into_iter() - .map(|node| node.weight) + .map(|node| node.weight.dist) } /// Return the [`Diagnostic`]s that were encountered while building the graph. @@ -352,7 +435,7 @@ impl ResolutionGraph { /// Return the underlying graph. pub fn petgraph( &self, - ) -> &petgraph::graph::Graph, petgraph::Directed> { + ) -> &petgraph::graph::Graph, petgraph::Directed> { &self.petgraph } @@ -506,7 +589,7 @@ impl ResolutionGraph { let mut locked_dists = vec![]; for node_index in self.petgraph.node_indices() { let dist = &self.petgraph[node_index]; - let mut locked_dist = lock::Distribution::from_resolved_dist(dist)?; + let mut locked_dist = lock::Distribution::from_annotated_dist(dist)?; for edge in self.petgraph.neighbors(node_index) { let dependency_dist = &self.petgraph[edge]; locked_dist.add_dependency(dependency_dist); @@ -586,9 +669,9 @@ impl<'a> DisplayResolutionGraph<'a> { #[derive(Debug)] enum Node<'a> { /// A node linked to an editable distribution. - Editable(&'a PackageName, &'a LocalEditable), + Editable(&'a LocalEditable), /// A node linked to a non-editable distribution. - Distribution(&'a PackageName, &'a ResolvedDist, &'a [ExtraName]), + Distribution(&'a PackageName, &'a AnnotatedDist, &'a [ExtraName]), } #[derive(Debug, PartialEq, Eq, PartialOrd, Ord)] @@ -600,18 +683,10 @@ enum NodeKey<'a> { } impl<'a> Node<'a> { - /// Return the name of the package. - fn name(&self) -> &'a PackageName { - match self { - Node::Editable(name, _) => name, - Node::Distribution(name, _, _) => name, - } - } - /// Return a comparable key for the node. fn key(&self) -> NodeKey<'a> { match self { - Node::Editable(_, editable) => NodeKey::Editable(editable.verbatim()), + Node::Editable(editable) => NodeKey::Editable(editable.verbatim()), Node::Distribution(name, _, _) => NodeKey::Distribution(name), } } @@ -619,8 +694,8 @@ impl<'a> Node<'a> { /// Return the [`IndexUrl`] of the distribution, if any. fn index(&self) -> Option<&IndexUrl> { match self { - Node::Editable(_, _) => None, - Node::Distribution(_, dist, _) => dist.index(), + Node::Editable(_) => None, + Node::Distribution(_, package, _) => package.dist.index(), } } } @@ -628,7 +703,7 @@ impl<'a> Node<'a> { impl Verbatim for Node<'_> { fn verbatim(&self) -> Cow<'_, str> { match self { - Node::Editable(_, editable) => Cow::Owned(format!("-e {}", editable.verbatim())), + Node::Editable(editable) => Cow::Owned(format!("-e {}", editable.verbatim())), Node::Distribution(_, dist, &[]) => dist.verbatim(), Node::Distribution(_, dist, extras) => { let mut extras = extras.to_vec(); @@ -661,7 +736,7 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> { } let node = if let Some((editable, _, _)) = self.resolution.editables.get(name) { - Node::Editable(name, editable) + Node::Editable(editable) } else if self.include_extras { Node::Distribution( name, @@ -689,13 +764,8 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> { // Display the distribution hashes, if any. let mut has_hashes = false; if self.show_hashes { - if let Some(hashes) = self - .resolution - .hashes - .get(node.name()) - .filter(|hashes| !hashes.is_empty()) - { - for hash in hashes { + if let Node::Distribution(_, package, _) = node { + for hash in &package.hashes { has_hashes = true; line.push_str(" \\\n"); line.push_str(" --hash="); @@ -722,7 +792,7 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> { // Include all external sources (e.g., requirements files). let default = BTreeSet::default(); let source = match node { - Node::Editable(_, editable) => { + Node::Editable(editable) => { self.sources.get_editable(&editable.url).unwrap_or(&default) } Node::Distribution(name, _, _) => self.sources.get(name).unwrap_or(&default), @@ -810,7 +880,7 @@ impl From for distribution_types::Resolution { .map(|node| { ( graph.petgraph[node].name().clone(), - graph.petgraph[node].clone(), + graph.petgraph[node].dist.clone(), ) }) .collect(), diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index 8ece382e2..29515defc 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -963,8 +963,8 @@ impl<'a, Provider: ResolverProvider, InstalledPackages: InstalledPackagesProvide // If we're excluding transitive dependencies, short-circuit. if self.dependency_mode.is_direct() { // If an extra is provided, wait for the metadata to be available, since it's - // still required for reporting diagnostics. - if extra.is_some() && !self.editables.contains(package_name) { + // still required for generating the lock file. + if !self.editables.contains(package_name) { // Determine the distribution to lookup. let dist = match url { Some(url) => PubGrubDistribution::from_url(package_name, url),