From 8001a90c86bb5759b96200386d5fe9e98a3dc677 Mon Sep 17 00:00:00 2001
From: Chan Kang <chankang17@gmail.com>
Date: Sat, 29 Jun 2024 16:54:03 -0400
Subject: [PATCH] fix the issue with pruning the last package

---
 crates/uv/src/commands/pip/tree.rs | 10 +++----
 crates/uv/tests/pip_tree.rs        | 43 ++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/crates/uv/src/commands/pip/tree.rs b/crates/uv/src/commands/pip/tree.rs
index 8b80213ee2b..a7001918772 100644
--- a/crates/uv/src/commands/pip/tree.rs
+++ b/crates/uv/src/commands/pip/tree.rs
@@ -165,11 +165,6 @@ impl<'a> DisplayDependencyGraph<'a> {
             return Vec::new();
         }
 
-        // Short-circuit if the current package is given in the prune list.
-        if self.prune.contains(installed_dist.name()) {
-            return Vec::new();
-        }
-
         let package_name = installed_dist.name().to_string();
         let is_visited = visited.contains(&package_name);
         let line = format!("{} v{}", package_name, installed_dist.version());
@@ -185,7 +180,10 @@ impl<'a> DisplayDependencyGraph<'a> {
 
         path.push(package_name.clone());
         visited.insert(package_name.clone());
-        let required_packages = required_with_no_extra(installed_dist, self.markers);
+        let required_packages = required_with_no_extra(installed_dist, self.markers)
+            .into_iter()
+            .filter(|p| !self.prune.contains(&p.name))
+            .collect::<Vec<_>>();
         for (index, required_package) in required_packages.iter().enumerate() {
             // Skip if the current package is not one of the installed distributions.
             if !self
diff --git a/crates/uv/tests/pip_tree.rs b/crates/uv/tests/pip_tree.rs
index 7b5a444b9e6..db9983ca412 100644
--- a/crates/uv/tests/pip_tree.rs
+++ b/crates/uv/tests/pip_tree.rs
@@ -31,6 +31,49 @@ fn no_package() {
     );
 }
 
+#[test]
+fn prune_last_in_the_subgroup() {
+    let context = TestContext::new("3.12");
+
+    let requirements_txt = context.temp_dir.child("requirements.txt");
+    requirements_txt.write_str("requests==2.31.0").unwrap();
+
+    uv_snapshot!(context
+        .pip_install()
+        .arg("-r")
+        .arg("requirements.txt")
+        .arg("--strict"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 5 packages in [TIME]
+    Prepared 5 packages in [TIME]
+    Installed 5 packages in [TIME]
+     + certifi==2024.2.2
+     + charset-normalizer==3.3.2
+     + idna==3.6
+     + requests==2.31.0
+     + urllib3==2.2.1
+    "###
+    );
+
+    context.assert_command("import requests").success();
+    uv_snapshot!(context.filters(), tree_command(&context).arg("--prune").arg("certifi"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    requests v2.31.0
+    ├── charset-normalizer v3.3.2
+    ├── idna v3.6
+    └── urllib3 v2.2.1
+
+    ----- stderr -----
+    "###
+    );
+}
+
 #[test]
 fn single_package() {
     let context = TestContext::new("3.12");