diff --git a/crates/uv/src/commands/pip/tree.rs b/crates/uv/src/commands/pip/tree.rs index 69a7412fa..d9b1f9c24 100644 --- a/crates/uv/src/commands/pip/tree.rs +++ b/crates/uv/src/commands/pip/tree.rs @@ -93,10 +93,10 @@ pub(crate) fn pip_tree( /// 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. /// This function will return `["charset-normalizer", "idna", "urllib", "certifi"]` for `requests`. -fn filtered_requirements( - dist: &InstalledDist, - markers: &MarkerEnvironment, -) -> Result>> { +fn filtered_requirements<'env>( + dist: &'env InstalledDist, + markers: &'env MarkerEnvironment, +) -> Result> + 'env> { Ok(dist .metadata()? .requires_dist @@ -106,51 +106,57 @@ fn filtered_requirements( .marker .as_ref() .map_or(true, |m| m.evaluate(markers, &[])) - }) - .collect::>()) + })) } #[derive(Debug)] -struct DisplayDependencyGraph<'a> { - site_packages: &'a SitePackages, +struct DisplayDependencyGraph<'env> { + // Installed packages. + site_packages: &'env SitePackages, /// 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 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, /// Whether to de-duplicate the displayed dependencies. no_dedupe: bool, - /// Map from package name to the list of required (reversed if --invert is given) packages. - requires_map: HashMap>, + /// Map from package name to its requirements. + /// + /// If `--invert` is given the map is inverted. + requirements: HashMap>, } -impl<'a> DisplayDependencyGraph<'a> { +impl<'env> DisplayDependencyGraph<'env> { /// Create a new [`DisplayDependencyGraph`] for the set of installed distributions. fn new( - site_packages: &'a SitePackages, + site_packages: &'env SitePackages, depth: usize, prune: Vec, no_dedupe: bool, invert: bool, - markers: &'a MarkerEnvironment, - ) -> Result> { - let mut dist_by_package_name = HashMap::new(); - let mut requires_map = HashMap::new(); + markers: &'env MarkerEnvironment, + ) -> Result> { + let mut distributions = HashMap::new(); + let mut requirements: HashMap<_, Vec<_>> = HashMap::new(); + + // Add all installed distributions. 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 required in filtered_requirements(site_package, markers)? { if invert { - requires_map + requirements .entry(required.name.clone()) - .or_insert_with(Vec::new) + .or_default() .push(site_package.name().clone()); } else { - requires_map + requirements .entry(site_package.name().clone()) - .or_insert_with(Vec::new) + .or_default() .push(required.name.clone()); } } @@ -158,20 +164,20 @@ impl<'a> DisplayDependencyGraph<'a> { Ok(Self { site_packages, - dist_by_package_name, + distributions, depth, prune, no_dedupe, - requires_map, + requirements, }) } /// Perform a depth-first traversal of the given distribution and its dependencies. fn visit( &self, - installed_dist: &InstalledDist, - visited: &mut FxHashMap>, - path: &mut Vec, + installed_dist: &'env InstalledDist, + visited: &mut FxHashMap<&'env PackageName, Vec>, + path: &mut Vec<&'env PackageName>, ) -> Result> { // Short-circuit if the current path is longer than the provided 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). // 2. The package has been visited and de-duplication is enabled (default). 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() { vec![line] } else { @@ -194,21 +200,24 @@ impl<'a> DisplayDependencyGraph<'a> { } } - let requirements_before_filtering = self.requires_map.get(installed_dist.name()); - let requirements = match requirements_before_filtering { - Some(requirements) => requirements - .iter() - .filter(|req| { - // Skip if the current package is not one of the installed distributions. - !self.prune.contains(req) && self.dist_by_package_name.contains_key(req) - }) - .cloned() - .collect(), - None => Vec::new(), - }; + let requirements = self + .requirements + .get(installed_dist.name()) + .into_iter() + .flatten() + .filter(|req| { + // Skip if the current package is not one of the installed distributions. + !self.prune.contains(req) && self.distributions.contains_key(req) + }) + .cloned() + .collect::>(); + 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 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 @@ -237,19 +246,17 @@ impl<'a> DisplayDependencyGraph<'a> { let mut prefixed_lines = Vec::new(); for (visited_index, visited_line) in self - .visit(self.dist_by_package_name[req], visited, path)? + .visit(self.distributions[req], visited, path)? .iter() .enumerate() { - prefixed_lines.push(format!( - "{}{}", - if visited_index == 0 { - prefix_top - } else { - prefix_rest - }, - visited_line - )); + let prefix = if visited_index == 0 { + prefix_top + } else { + prefix_rest + }; + + prefixed_lines.push(format!("{prefix}{visited_line}")); } lines.extend(prefixed_lines); } @@ -260,19 +267,17 @@ impl<'a> DisplayDependencyGraph<'a> { /// Depth-first traverse the nodes to render the tree. fn render(&self) -> Result> { - let mut visited: FxHashMap> = FxHashMap::default(); - let mut path: Vec = Vec::new(); + let mut visited: FxHashMap<&PackageName, Vec> = FxHashMap::default(); + let mut path: Vec<&PackageName> = Vec::new(); let mut lines: Vec = Vec::new(); - // The starting nodes are those that are not required by any other package. - let mut non_starting_nodes = HashSet::new(); - for children in self.requires_map.values() { - non_starting_nodes.extend(children); - } + // The root nodes are those that are not required by any other package. + let children: HashSet<_> = self.requirements.values().flatten().collect(); for site_package in self.site_packages.iter() { // If the current package is not required by any other package, start the traversal // 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)?); } }