mirror of
				https://github.com/astral-sh/uv.git
				synced 2025-10-31 20:09:09 +00:00 
			
		
		
		
	Create distinct graph nodes for each package extra (#3908)
## Summary Today, we represent each package as a single node in the graph, and combine all the extras. This is helpful for the `requirements.txt`-style resolution, in which we want to show each a single line for each package with the extras combined into a single array. This PR modifies the representation to instead use a separate node for each (package, extra) pair. We then reduce into the previous format when printing in the `requirements.txt`-style format, so there shouldn't be any user-facing changes here.
This commit is contained in:
		
							parent
							
								
									0edb660205
								
							
						
					
					
						commit
						19c91e7dac
					
				
					 2 changed files with 97 additions and 110 deletions
				
			
		|  | @ -1,12 +1,15 @@ | |||
| use std::collections::BTreeSet; | ||||
| use std::hash::BuildHasherDefault; | ||||
| 
 | ||||
| use owo_colors::OwoColorize; | ||||
| use petgraph::visit::EdgeRef; | ||||
| use petgraph::Direction; | ||||
| use rustc_hash::FxHashMap; | ||||
| 
 | ||||
| use distribution_types::{Name, SourceAnnotations}; | ||||
| use uv_normalize::PackageName; | ||||
| 
 | ||||
| use crate::resolution::AnnotatedDist; | ||||
| use crate::ResolutionGraph; | ||||
| 
 | ||||
| /// A [`std::fmt::Display`] implementation for the resolution graph.
 | ||||
|  | @ -77,13 +80,53 @@ impl<'a> DisplayResolutionGraph<'a> { | |||
| /// Write the graph in the `{name}=={version}` format of requirements.txt that pip uses.
 | ||||
| impl std::fmt::Display for DisplayResolutionGraph<'_> { | ||||
|     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { | ||||
|         // Reduce the graph, such that all nodes for a single package are combined, regardless of
 | ||||
|         // the extras.
 | ||||
|         //
 | ||||
|         // For example, `flask` and `flask[dotenv]` should be reduced into a single `flask[dotenv]`
 | ||||
|         // node.
 | ||||
|         let petgraph = { | ||||
|             let mut petgraph = | ||||
|                 petgraph::graph::Graph::<AnnotatedDist, (), petgraph::Directed>::with_capacity( | ||||
|                     self.resolution.petgraph.node_count(), | ||||
|                     self.resolution.petgraph.edge_count(), | ||||
|                 ); | ||||
|             let mut inverse = FxHashMap::with_capacity_and_hasher( | ||||
|                 self.resolution.petgraph.node_count(), | ||||
|                 BuildHasherDefault::default(), | ||||
|             ); | ||||
| 
 | ||||
|             // Re-add the nodes to the reduced graph.
 | ||||
|             for index in self.resolution.petgraph.node_indices() { | ||||
|                 let dist = &self.resolution.petgraph[index]; | ||||
| 
 | ||||
|                 if let Some(index) = inverse.get(dist.name()) { | ||||
|                     let node: &mut AnnotatedDist = &mut petgraph[*index]; | ||||
|                     node.extras.extend(dist.extras.iter().cloned()); | ||||
|                     node.extras.sort_unstable(); | ||||
|                     node.extras.dedup(); | ||||
|                 } else { | ||||
|                     let index = petgraph.add_node(dist.clone()); | ||||
|                     inverse.insert(dist.name(), index); | ||||
|                 } | ||||
|             } | ||||
| 
 | ||||
|             // Re-add the edges to the reduced graph.
 | ||||
|             for edge in self.resolution.petgraph.edge_indices() { | ||||
|                 let (source, target) = self.resolution.petgraph.edge_endpoints(edge).unwrap(); | ||||
|                 let source = inverse[self.resolution.petgraph[source].name()]; | ||||
|                 let target = inverse[self.resolution.petgraph[target].name()]; | ||||
|                 petgraph.update_edge(source, target, ()); | ||||
|             } | ||||
| 
 | ||||
|             petgraph | ||||
|         }; | ||||
| 
 | ||||
|         // Collect all packages.
 | ||||
|         let mut nodes = self | ||||
|             .resolution | ||||
|             .petgraph | ||||
|         let mut nodes = petgraph | ||||
|             .node_indices() | ||||
|             .filter_map(|index| { | ||||
|                 let dist = &self.resolution.petgraph[index]; | ||||
|                 let dist = &petgraph[index]; | ||||
|                 let name = dist.name(); | ||||
|                 if self.no_emit_packages.contains(name) { | ||||
|                     return None; | ||||
|  | @ -119,11 +162,9 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> { | |||
|             // package (e.g., `# via mypy`).
 | ||||
|             if self.include_annotations { | ||||
|                 // Display all dependencies.
 | ||||
|                 let mut edges = self | ||||
|                     .resolution | ||||
|                     .petgraph | ||||
|                 let mut edges = petgraph | ||||
|                     .edges_directed(index, Direction::Incoming) | ||||
|                     .map(|edge| &self.resolution.petgraph[edge.source()]) | ||||
|                     .map(|edge| &petgraph[edge.source()]) | ||||
|                     .collect::<Vec<_>>(); | ||||
|                 edges.sort_unstable_by_key(|package| package.name()); | ||||
| 
 | ||||
|  |  | |||
|  | @ -46,102 +46,17 @@ impl ResolutionGraph { | |||
|         state: &State<UvDependencyProvider>, | ||||
|         preferences: &Preferences, | ||||
|     ) -> anyhow::Result<Self, ResolveError> { | ||||
|         // Collect and validate the extras.
 | ||||
|         let mut extras = FxHashMap::default(); | ||||
|         let mut diagnostics = Vec::new(); | ||||
|         for (package, version) in selection { | ||||
|             match &**package { | ||||
|                 PubGrubPackageInner::Package { | ||||
|                     name, | ||||
|                     extra: Some(extra), | ||||
|                     marker: None, | ||||
|                     url: None, | ||||
|                 } => { | ||||
|                     let dist = PubGrubDistribution::from_registry(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() | ||||
|                         ) | ||||
|                     }; | ||||
| 
 | ||||
|                     if archive.metadata.provides_extras.contains(extra) { | ||||
|                         extras | ||||
|                             .entry(name.clone()) | ||||
|                             .or_insert_with(Vec::new) | ||||
|                             .push(extra.clone()); | ||||
|                     } else { | ||||
|                         let dist = pins | ||||
|                             .get(name, version) | ||||
|                             .unwrap_or_else(|| panic!("Every package should be pinned: {name:?}")) | ||||
|                             .clone(); | ||||
| 
 | ||||
|                         diagnostics.push(ResolutionDiagnostic::MissingExtra { | ||||
|                             dist, | ||||
|                             extra: extra.clone(), | ||||
|                         }); | ||||
|                     } | ||||
|                 } | ||||
|                 PubGrubPackageInner::Package { | ||||
|                     name, | ||||
|                     extra: Some(extra), | ||||
|                     marker: None, | ||||
|                     url: Some(url), | ||||
|                 } => { | ||||
|                     let dist = PubGrubDistribution::from_url(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() | ||||
|                         ) | ||||
|                     }; | ||||
| 
 | ||||
|                     if archive.metadata.provides_extras.contains(extra) { | ||||
|                         extras | ||||
|                             .entry(name.clone()) | ||||
|                             .or_insert_with(Vec::new) | ||||
|                             .push(extra.clone()); | ||||
|                     } else { | ||||
|                         let dist = Dist::from_url(name.clone(), url_to_precise(url.clone()))?; | ||||
| 
 | ||||
|                         diagnostics.push(ResolutionDiagnostic::MissingExtra { | ||||
|                             dist: dist.into(), | ||||
|                             extra: extra.clone(), | ||||
|                         }); | ||||
|                     } | ||||
|                 } | ||||
|                 _ => {} | ||||
|             }; | ||||
|         } | ||||
| 
 | ||||
|         // 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 inverse = | ||||
|             FxHashMap::with_capacity_and_hasher(selection.len(), BuildHasherDefault::default()); | ||||
|         let mut diagnostics = Vec::new(); | ||||
| 
 | ||||
|         for (package, version) in selection { | ||||
|             match &**package { | ||||
|                 PubGrubPackageInner::Package { | ||||
|                     name, | ||||
|                     extra: None, | ||||
|                     extra, | ||||
|                     marker: None, | ||||
|                     url: None, | ||||
|                 } => { | ||||
|  | @ -213,8 +128,18 @@ impl ResolutionGraph { | |||
|                         archive.metadata.clone() | ||||
|                     }; | ||||
| 
 | ||||
|                     // Validate the extra.
 | ||||
|                     if let Some(extra) = extra { | ||||
|                         if !metadata.provides_extras.contains(extra) { | ||||
|                             diagnostics.push(ResolutionDiagnostic::MissingExtra { | ||||
|                                 dist: dist.clone(), | ||||
|                                 extra: extra.clone(), | ||||
|                             }); | ||||
|                         } | ||||
|                     } | ||||
| 
 | ||||
|                     // Extract the extras.
 | ||||
|                     let extras = extras.get(name).cloned().unwrap_or_default(); | ||||
|                     let extras = extra.clone().map(|extra| vec![extra]).unwrap_or_default(); | ||||
| 
 | ||||
|                     // Add the distribution to the graph.
 | ||||
|                     let index = petgraph.add_node(AnnotatedDist { | ||||
|  | @ -223,16 +148,16 @@ impl ResolutionGraph { | |||
|                         hashes, | ||||
|                         metadata, | ||||
|                     }); | ||||
|                     inverse.insert(name, index); | ||||
|                     inverse.insert((name, extra), index); | ||||
|                 } | ||||
| 
 | ||||
|                 PubGrubPackageInner::Package { | ||||
|                     name, | ||||
|                     extra: None, | ||||
|                     extra, | ||||
|                     marker: None, | ||||
|                     url: Some(url), | ||||
|                 } => { | ||||
|                     // Create the distribution.
 | ||||
| 
 | ||||
|                     let dist = Dist::from_url(name.clone(), url_to_precise(url.clone()))?; | ||||
| 
 | ||||
|                     // Extract the hashes, preserving those that were already present in the
 | ||||
|  | @ -275,8 +200,18 @@ impl ResolutionGraph { | |||
|                         archive.metadata.clone() | ||||
|                     }; | ||||
| 
 | ||||
|                     // Validate the extra.
 | ||||
|                     if let Some(extra) = extra { | ||||
|                         if !metadata.provides_extras.contains(extra) { | ||||
|                             diagnostics.push(ResolutionDiagnostic::MissingExtra { | ||||
|                                 dist: dist.clone().into(), | ||||
|                                 extra: extra.clone(), | ||||
|                             }); | ||||
|                         } | ||||
|                     } | ||||
| 
 | ||||
|                     // Extract the extras.
 | ||||
|                     let extras = extras.get(name).cloned().unwrap_or_default(); | ||||
|                     let extras = extra.clone().map(|extra| vec![extra]).unwrap_or_default(); | ||||
| 
 | ||||
|                     // Add the distribution to the graph.
 | ||||
|                     let index = petgraph.add_node(AnnotatedDist { | ||||
|  | @ -285,8 +220,9 @@ impl ResolutionGraph { | |||
|                         hashes, | ||||
|                         metadata, | ||||
|                     }); | ||||
|                     inverse.insert(name, index); | ||||
|                     inverse.insert((name, extra), index); | ||||
|                 } | ||||
| 
 | ||||
|                 _ => {} | ||||
|             }; | ||||
|         } | ||||
|  | @ -305,14 +241,22 @@ impl ResolutionGraph { | |||
|                         continue; | ||||
|                     } | ||||
| 
 | ||||
|                     if !self_version.contains(version) { | ||||
|                         continue; | ||||
|                     } | ||||
| 
 | ||||
|                     let PubGrubPackageInner::Package { | ||||
|                         name: self_name, .. | ||||
|                         name: self_name, | ||||
|                         extra: self_extra, | ||||
|                         .. | ||||
|                     } = &**self_package | ||||
|                     else { | ||||
|                         continue; | ||||
|                     }; | ||||
| 
 | ||||
|                     let PubGrubPackageInner::Package { | ||||
|                         name: dependency_name, | ||||
|                         extra: dependency_extra, | ||||
|                         .. | ||||
|                     } = &**dependency_package | ||||
|                     else { | ||||
|  | @ -324,11 +268,9 @@ impl ResolutionGraph { | |||
|                         continue; | ||||
|                     } | ||||
| 
 | ||||
|                     if self_version.contains(version) { | ||||
|                         let self_index = &inverse[self_name]; | ||||
|                         let dependency_index = &inverse[dependency_name]; | ||||
|                         petgraph.update_edge(*self_index, *dependency_index, ()); | ||||
|                     } | ||||
|                     let self_index = &inverse[&(self_name, self_extra)]; | ||||
|                     let dependency_index = &inverse[&(dependency_name, dependency_extra)]; | ||||
|                     petgraph.update_edge(*self_index, *dependency_index, ()); | ||||
|                 } | ||||
|             } | ||||
|         } | ||||
|  | @ -339,14 +281,18 @@ impl ResolutionGraph { | |||
|         }) | ||||
|     } | ||||
| 
 | ||||
|     /// Return the number of packages in the graph.
 | ||||
|     /// Return the number of distinct packages in the graph.
 | ||||
|     pub fn len(&self) -> usize { | ||||
|         self.petgraph.node_count() | ||||
|         self.petgraph | ||||
|             .node_indices() | ||||
|             .map(|index| &self.petgraph[index]) | ||||
|             .filter(|dist| dist.extras.is_empty()) | ||||
|             .count() | ||||
|     } | ||||
| 
 | ||||
|     /// Return `true` if there are no packages in the graph.
 | ||||
|     pub fn is_empty(&self) -> bool { | ||||
|         self.petgraph.node_count() == 0 | ||||
|         self.len() == 0 | ||||
|     } | ||||
| 
 | ||||
|     /// Returns `true` if the graph contains the given package.
 | ||||
|  |  | |||
		Loading…
	
	Add table
		Add a link
		
	
		Reference in a new issue
	
	 Charlie Marsh
						Charlie Marsh