From 8a2af8bc8349554ee134c6487f4b20e21f1a046b Mon Sep 17 00:00:00 2001 From: Ibraheem Ahmed Date: Mon, 1 Jul 2024 15:21:38 -0400 Subject: [PATCH] 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 | 14 ++++--- crates/uv/src/commands/pip/tree.rs | 35 ++++++----------- crates/uv/tests/common/mod.rs | 8 ++++ crates/uv/tests/pip_tree.rs | 49 ++++++++++-------------- 4 files changed, 50 insertions(+), 56 deletions(-) diff --git a/crates/uv-installer/src/site_packages.rs b/crates/uv-installer/src/site_packages.rs index a116c2182..278b6b056 100644 --- a/crates/uv-installer/src/site_packages.rs +++ b/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 - .entry(dist.url.clone()) - .or_insert_with(Vec::new) - .push(idx); + by_url.entry(dist.url.clone()).or_default().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 { let Some(indexes) = self.by_name.get(name) else { diff --git a/crates/uv/src/commands/pip/tree.rs b/crates/uv/src/commands/pip/tree.rs index d9b1f9c24..f1e70cdb7 100644 --- a/crates/uv/src/commands/pip/tree.rs +++ b/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> { - 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::>(); @@ -244,21 +235,19 @@ impl<'env> DisplayDependencyGraph<'env> { ("├── ", "│ ") }; - let mut prefixed_lines = Vec::new(); - for (visited_index, visited_line) in self - .visit(self.distributions[req], visited, path)? - .iter() - .enumerate() - { - let prefix = if visited_index == 0 { - prefix_top - } else { - prefix_rest - }; + for distribution in self.site_packages.get_packages(req) { + for (visited_index, visited_line) in + self.visit(distribution, visited, path)?.iter().enumerate() + { + let prefix = if visited_index == 0 { + prefix_top + } else { + prefix_rest + }; - prefixed_lines.push(format!("{prefix}{visited_line}")); + lines.push(format!("{prefix}{visited_line}")); + } } - lines.extend(prefixed_lines); } path.pop(); diff --git a/crates/uv/tests/common/mod.rs b/crates/uv/tests/common/mod.rs index c78d6cc88..1bce8bfb7 100644 --- a/crates/uv/tests/common/mod.rs +++ b/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()); diff --git a/crates/uv/tests/pip_tree.rs b/crates/uv/tests/pip_tree.rs index 71c798725..c490b3234 100644 --- a/crates/uv/tests/pip_tree.rs +++ b/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::>(); - uv_snapshot!(filters, tree_command(&context), @r###" + uv_snapshot!(filters, context.pip_tree(), @r###" success: true exit_code: 0 ----- stdout -----