Clean up uv pip tree code (#4700)

## Summary

Minor improvements/nits.
This commit is contained in:
Ibraheem Ahmed 2024-07-01 13:56:16 -04:00 committed by GitHub
parent 58499439d3
commit 348efa26ba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -93,10 +93,10 @@ pub(crate) fn pip_tree(
/// For example, `requests==2.32.3` requires `charset-normalizer`, `idna`, `urllib`, and `certifi` at /// For example, `requests==2.32.3` requires `charset-normalizer`, `idna`, `urllib`, and `certifi` at
/// all times, `PySocks` on `socks` extra and `chardet` on `use_chardet_on_py3` extra. /// all times, `PySocks` on `socks` extra and `chardet` on `use_chardet_on_py3` extra.
/// This function will return `["charset-normalizer", "idna", "urllib", "certifi"]` for `requests`. /// This function will return `["charset-normalizer", "idna", "urllib", "certifi"]` for `requests`.
fn filtered_requirements( fn filtered_requirements<'env>(
dist: &InstalledDist, dist: &'env InstalledDist,
markers: &MarkerEnvironment, markers: &'env MarkerEnvironment,
) -> Result<Vec<Requirement<VerbatimParsedUrl>>> { ) -> Result<impl Iterator<Item = Requirement<VerbatimParsedUrl>> + 'env> {
Ok(dist Ok(dist
.metadata()? .metadata()?
.requires_dist .requires_dist
@ -106,51 +106,57 @@ fn filtered_requirements(
.marker .marker
.as_ref() .as_ref()
.map_or(true, |m| m.evaluate(markers, &[])) .map_or(true, |m| m.evaluate(markers, &[]))
}) }))
.collect::<Vec<_>>())
} }
#[derive(Debug)] #[derive(Debug)]
struct DisplayDependencyGraph<'a> { struct DisplayDependencyGraph<'env> {
site_packages: &'a SitePackages, // Installed packages.
site_packages: &'env SitePackages,
/// Map from package name to the installed distribution. /// Map from package name to the installed distribution.
dist_by_package_name: HashMap<&'a PackageName, &'a InstalledDist>, distributions: HashMap<&'env PackageName, &'env InstalledDist>,
/// Maximum display depth of the dependency tree /// Maximum display depth of the dependency tree
depth: usize, depth: usize,
/// Prune the given package from the display of the dependency tree. /// Prune the given packages from the display of the dependency tree.
prune: Vec<PackageName>, prune: Vec<PackageName>,
/// Whether to de-duplicate the displayed dependencies. /// Whether to de-duplicate the displayed dependencies.
no_dedupe: bool, no_dedupe: bool,
/// Map from package name to the list of required (reversed if --invert is given) packages. /// Map from package name to its requirements.
requires_map: HashMap<PackageName, Vec<PackageName>>, ///
/// If `--invert` is given the map is inverted.
requirements: HashMap<PackageName, Vec<PackageName>>,
} }
impl<'a> DisplayDependencyGraph<'a> { impl<'env> DisplayDependencyGraph<'env> {
/// Create a new [`DisplayDependencyGraph`] for the set of installed distributions. /// Create a new [`DisplayDependencyGraph`] for the set of installed distributions.
fn new( fn new(
site_packages: &'a SitePackages, site_packages: &'env SitePackages,
depth: usize, depth: usize,
prune: Vec<PackageName>, prune: Vec<PackageName>,
no_dedupe: bool, no_dedupe: bool,
invert: bool, invert: bool,
markers: &'a MarkerEnvironment, markers: &'env MarkerEnvironment,
) -> Result<DisplayDependencyGraph<'a>> { ) -> Result<DisplayDependencyGraph<'env>> {
let mut dist_by_package_name = HashMap::new(); let mut distributions = HashMap::new();
let mut requires_map = HashMap::new(); let mut requirements: HashMap<_, Vec<_>> = HashMap::new();
// Add all installed distributions.
for site_package in site_packages.iter() { for site_package in site_packages.iter() {
dist_by_package_name.insert(site_package.name(), site_package); distributions.insert(site_package.name(), site_package);
} }
// Add all transitive requirements.
for site_package in site_packages.iter() { for site_package in site_packages.iter() {
for required in filtered_requirements(site_package, markers)? { for required in filtered_requirements(site_package, markers)? {
if invert { if invert {
requires_map requirements
.entry(required.name.clone()) .entry(required.name.clone())
.or_insert_with(Vec::new) .or_default()
.push(site_package.name().clone()); .push(site_package.name().clone());
} else { } else {
requires_map requirements
.entry(site_package.name().clone()) .entry(site_package.name().clone())
.or_insert_with(Vec::new) .or_default()
.push(required.name.clone()); .push(required.name.clone());
} }
} }
@ -158,20 +164,20 @@ impl<'a> DisplayDependencyGraph<'a> {
Ok(Self { Ok(Self {
site_packages, site_packages,
dist_by_package_name, distributions,
depth, depth,
prune, prune,
no_dedupe, no_dedupe,
requires_map, requirements,
}) })
} }
/// Perform a depth-first traversal of the given distribution and its dependencies. /// Perform a depth-first traversal of the given distribution and its dependencies.
fn visit( fn visit(
&self, &self,
installed_dist: &InstalledDist, installed_dist: &'env InstalledDist,
visited: &mut FxHashMap<PackageName, Vec<PackageName>>, visited: &mut FxHashMap<&'env PackageName, Vec<PackageName>>,
path: &mut Vec<PackageName>, path: &mut Vec<&'env PackageName>,
) -> Result<Vec<String>> { ) -> Result<Vec<String>> {
// Short-circuit if the current path is longer than the provided depth. // Short-circuit if the current path is longer than the provided depth.
if path.len() > self.depth { if path.len() > self.depth {
@ -185,7 +191,7 @@ impl<'a> DisplayDependencyGraph<'a> {
// 1. The package is in the current traversal path (i.e., a dependency cycle). // 1. The package is in the current traversal path (i.e., a dependency cycle).
// 2. The package has been visited and de-duplication is enabled (default). // 2. The package has been visited and de-duplication is enabled (default).
if let Some(requirements) = visited.get(package_name) { if let Some(requirements) = visited.get(package_name) {
if !self.no_dedupe || path.contains(package_name) { if !self.no_dedupe || path.contains(&package_name) {
return Ok(if requirements.is_empty() { return Ok(if requirements.is_empty() {
vec![line] vec![line]
} else { } else {
@ -194,21 +200,24 @@ impl<'a> DisplayDependencyGraph<'a> {
} }
} }
let requirements_before_filtering = self.requires_map.get(installed_dist.name()); let requirements = self
let requirements = match requirements_before_filtering { .requirements
Some(requirements) => requirements .get(installed_dist.name())
.iter() .into_iter()
.filter(|req| { .flatten()
// Skip if the current package is not one of the installed distributions. .filter(|req| {
!self.prune.contains(req) && self.dist_by_package_name.contains_key(req) // Skip if the current package is not one of the installed distributions.
}) !self.prune.contains(req) && self.distributions.contains_key(req)
.cloned() })
.collect(), .cloned()
None => Vec::new(), .collect::<Vec<_>>();
};
let mut lines = vec![line]; let mut lines = vec![line];
visited.insert(package_name.clone(), requirements.clone());
path.push(package_name.clone()); // Keep track of the dependency path to avoid cycles.
visited.insert(package_name, requirements.clone());
path.push(package_name);
for (index, req) in requirements.iter().enumerate() { for (index, req) in requirements.iter().enumerate() {
// For sub-visited packages, add the prefix to make the tree display user-friendly. // For sub-visited packages, add the prefix to make the tree display user-friendly.
// The key observation here is you can group the tree as follows when you're at the // The key observation here is you can group the tree as follows when you're at the
@ -237,19 +246,17 @@ impl<'a> DisplayDependencyGraph<'a> {
let mut prefixed_lines = Vec::new(); let mut prefixed_lines = Vec::new();
for (visited_index, visited_line) in self for (visited_index, visited_line) in self
.visit(self.dist_by_package_name[req], visited, path)? .visit(self.distributions[req], visited, path)?
.iter() .iter()
.enumerate() .enumerate()
{ {
prefixed_lines.push(format!( let prefix = if visited_index == 0 {
"{}{}", prefix_top
if visited_index == 0 { } else {
prefix_top prefix_rest
} else { };
prefix_rest
}, prefixed_lines.push(format!("{prefix}{visited_line}"));
visited_line
));
} }
lines.extend(prefixed_lines); lines.extend(prefixed_lines);
} }
@ -260,19 +267,17 @@ impl<'a> DisplayDependencyGraph<'a> {
/// Depth-first traverse the nodes to render the tree. /// Depth-first traverse the nodes to render the tree.
fn render(&self) -> Result<Vec<String>> { fn render(&self) -> Result<Vec<String>> {
let mut visited: FxHashMap<PackageName, Vec<PackageName>> = FxHashMap::default(); let mut visited: FxHashMap<&PackageName, Vec<PackageName>> = FxHashMap::default();
let mut path: Vec<PackageName> = Vec::new(); let mut path: Vec<&PackageName> = Vec::new();
let mut lines: Vec<String> = Vec::new(); let mut lines: Vec<String> = Vec::new();
// The starting nodes are those that are not required by any other package. // The root nodes are those that are not required by any other package.
let mut non_starting_nodes = HashSet::new(); let children: HashSet<_> = self.requirements.values().flatten().collect();
for children in self.requires_map.values() {
non_starting_nodes.extend(children);
}
for site_package in self.site_packages.iter() { for site_package in self.site_packages.iter() {
// If the current package is not required by any other package, start the traversal // If the current package is not required by any other package, start the traversal
// with the current package as the root. // with the current package as the root.
if !non_starting_nodes.contains(site_package.name()) { if !children.contains(site_package.name()) {
path.clear();
lines.extend(self.visit(site_package, &mut visited, &mut path)?); lines.extend(self.visit(site_package, &mut visited, &mut path)?);
} }
} }