More uv pip tree cleanup (#4702)

## Summary

Follow up to https://github.com/astral-sh/uv/pull/4700.

crates/uv-installer/src/site_packages.rs

@@ -87,15 +87,12 @@ impl SitePackages {
                 // Index the distribution by name.
                 by_name
                     .entry(dist_info.name().clone())
-                    .or_insert_with(Vec::new)
+                    .or_default()
                     .push(idx);
 
                 // Index the distribution by URL.
                 if let InstalledDist::Url(dist) = &dist_info {
-                    by_url
+                    by_url.entry(dist.url.clone()).or_default().push(idx);
-                        .entry(dist.url.clone())
-                        .or_insert_with(Vec::new)
-                        .push(idx);
                 }
 
                 // Add the distribution to the database.
@@ -127,6 +124,13 @@ impl SitePackages {
             .collect()
     }
 
+    /// Returns `true` if there are any installed distributions with the given package name.
+    pub fn contains_package(&self, name: &PackageName) -> bool {
+        self.by_name
+            .get(name)
+            .is_some_and(|packages| !packages.is_empty())
+    }
+
     /// Remove the given packages from the index, returning all installed versions, if any.
     pub fn remove_packages(&mut self, name: &PackageName) -> Vec<InstalledDist> {
         let Some(indexes) = self.by_name.get(name) else {

crates/uv/src/commands/pip/tree.rs

@@ -113,8 +113,6 @@ fn filtered_requirements<'env>(
 struct DisplayDependencyGraph<'env> {
     // Installed packages.
     site_packages: &'env SitePackages,
-    /// Map from package name to the installed distribution.
-    distributions: HashMap<&'env PackageName, &'env InstalledDist>,
     /// Maximum display depth of the dependency tree
     depth: usize,
     /// Prune the given packages from the display of the dependency tree.
@@ -137,14 +135,8 @@ impl<'env> DisplayDependencyGraph<'env> {
         invert: bool,
         markers: &'env MarkerEnvironment,
     ) -> Result<DisplayDependencyGraph<'env>> {
-        let mut distributions = HashMap::new();
         let mut requirements: HashMap<_, Vec<_>> = HashMap::new();
 
-        // Add all installed distributions.
-        for site_package in site_packages.iter() {
-            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)? {
@@ -164,7 +156,6 @@ impl<'env> DisplayDependencyGraph<'env> {
 
         Ok(Self {
             site_packages,
-            distributions,
             depth,
             prune,
             no_dedupe,
@@ -207,7 +198,7 @@ impl<'env> DisplayDependencyGraph<'env> {
             .flatten()
             .filter(|req| {
                 // Skip if the current package is not one of the installed distributions.
-                !self.prune.contains(req) && self.distributions.contains_key(req)
+                !self.prune.contains(req) && self.site_packages.contains_package(req)
             })
             .cloned()
             .collect::<Vec<_>>();
@@ -244,21 +235,19 @@ impl<'env> DisplayDependencyGraph<'env> {
                 ("├── ", "│   ")
             };
 
-            let mut prefixed_lines = Vec::new();
+            for distribution in self.site_packages.get_packages(req) {
-            for (visited_index, visited_line) in self
+                for (visited_index, visited_line) in
-                .visit(self.distributions[req], visited, path)?
+                    self.visit(distribution, visited, path)?.iter().enumerate()
-                .iter()
+                {
-                .enumerate()
+                    let prefix = if visited_index == 0 {
-            {
+                        prefix_top
-                let prefix = if visited_index == 0 {
+                    } else {
-                    prefix_top
+                        prefix_rest
-                } else {
+                    };
-                    prefix_rest
-                };
 
-                prefixed_lines.push(format!("{prefix}{visited_line}"));
+                    lines.push(format!("{prefix}{visited_line}"));
+                }
             }
-            lines.extend(prefixed_lines);
         }
         path.pop();
 

crates/uv/tests/common/mod.rs

@@ -347,6 +347,14 @@ impl TestContext {
         command
     }
 
+    /// Create a `pip tree` command for testing.
+    pub fn pip_tree(&self) -> Command {
+        let mut command = Command::new(get_bin());
+        command.arg("pip").arg("tree");
+        self.add_shared_args(&mut command);
+        command
+    }
+
     /// Create a `uv sync` command with options shared across scenarios.
     pub fn sync(&self) -> Command {
         let mut command = Command::new(get_bin());

crates/uv/tests/pip_tree.rs

@@ -9,18 +9,11 @@ use crate::common::{get_bin, TestContext};
 
 mod common;
 
-fn tree_command(context: &TestContext) -> Command {
-    let mut command = Command::new(get_bin());
-    command.arg("pip").arg("tree");
-    context.add_shared_args(&mut command);
-    command
-}
-
 #[test]
 fn no_package() {
     let context = TestContext::new("3.12");
 
-    uv_snapshot!(context.filters(), tree_command(&context), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree(), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -60,7 +53,7 @@ fn prune_last_in_the_subgroup() {
     );
 
     context.assert_command("import requests").success();
-    uv_snapshot!(context.filters(), tree_command(&context).arg("--prune").arg("certifi"), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree().arg("--prune").arg("certifi"), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -103,7 +96,7 @@ fn single_package() {
     );
 
     context.assert_command("import requests").success();
-    uv_snapshot!(context.filters(), tree_command(&context), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree(), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -150,7 +143,7 @@ fn python_version_marker() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree(), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -196,7 +189,7 @@ fn nested_dependencies() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree(), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -212,7 +205,7 @@ fn nested_dependencies() {
     );
 }
 
-// identical test as `--invert` since `--reverse` is simply an alias for `--invert`.
+// Identical test as `invert` since `--reverse` is simply an alias for `--invert`.
 #[test]
 fn reverse() {
     let context = TestContext::new("3.12");
@@ -243,7 +236,7 @@ fn reverse() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context).arg("--reverse"), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree().arg("--reverse"), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -291,7 +284,7 @@ fn invert() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context).arg("--invert"), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree().arg("--invert"), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -509,7 +502,7 @@ fn prune() {
 
 #[test]
 #[cfg(target_os = "macos")]
-fn nested_dependencies_more_complex_inverted() {
+fn complex_nested_dependencies_inverted() {
     let context = TestContext::new("3.12");
 
     let requirements_txt = context.temp_dir.child("requirements.txt");
@@ -563,7 +556,7 @@ fn nested_dependencies_more_complex_inverted() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context).arg("--invert"), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree().arg("--invert"), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -631,7 +624,7 @@ fn nested_dependencies_more_complex_inverted() {
 
 #[test]
 #[cfg(target_os = "macos")]
-fn nested_dependencies_more_complex() {
+fn complex_nested_dependencies() {
     let context = TestContext::new("3.12");
 
     let requirements_txt = context.temp_dir.child("requirements.txt");
@@ -685,7 +678,7 @@ fn nested_dependencies_more_complex() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree(), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -734,7 +727,7 @@ fn nested_dependencies_more_complex() {
 
 #[test]
 #[cfg(target_os = "macos")]
-fn prune_big_tree() {
+fn prune_large_tree() {
     let context = TestContext::new("3.12");
 
     let requirements_txt = context.temp_dir.child("requirements.txt");
@@ -875,7 +868,7 @@ fn cyclic_dependency() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree(), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -932,7 +925,7 @@ fn removed_dependency() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree(), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -987,7 +980,7 @@ fn multiple_packages() {
         filters.push(("└── colorama v0.4.6\n", ""));
     }
     context.assert_command("import requests").success();
-    uv_snapshot!(filters, tree_command(&context), @r###"
+    uv_snapshot!(filters, context.pip_tree(), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -1046,7 +1039,7 @@ fn multiple_packages_shared_descendant() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree(), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -1114,7 +1107,7 @@ fn no_dedupe_and_invert() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context).arg("--no-dedupe").arg("--invert"), @r###"
+    uv_snapshot!(context.filters(), context.pip_tree().arg("--no-dedupe").arg("--invert"), @r###"
     success: true
     exit_code: 0
     ----- stdout -----
@@ -1212,7 +1205,7 @@ fn no_dedupe_and_cycle() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context)
+    uv_snapshot!(context.filters(), context.pip_tree()
         .arg("--no-dedupe"), @r###"
     success: true
     exit_code: 0
@@ -1290,7 +1283,7 @@ fn no_dedupe() {
     "###
     );
 
-    uv_snapshot!(context.filters(), tree_command(&context)
+    uv_snapshot!(context.filters(), context.pip_tree()
         .arg("--no-dedupe"), @r###"
     success: true
     exit_code: 0
@@ -1349,7 +1342,7 @@ fn with_editable() {
         .chain(vec![(r"\-\-\-\-\-\-+.*", "[UNDERLINE]"), ("  +", " ")])
         .collect::<Vec<_>>();
 
-    uv_snapshot!(filters, tree_command(&context), @r###"
+    uv_snapshot!(filters, context.pip_tree(), @r###"
     success: true
     exit_code: 0
     ----- stdout -----