Implement --package for pip tree (#4655)

## Summary

Part of https://github.com/astral-sh/uv/issues/4439.

## Test Plan

The existing tests pass + added a couple of tests to ensure `--package` behaves as expected.

crates/uv-cli/src/lib.rs

@@ -1428,6 +1428,11 @@ pub struct PipTreeArgs {
     /// Prune the given package from the display of the dependency tree.
     #[arg(long)]
     pub prune: Vec<PackageName>,
+
+    /// Display only the specified packages.
+    #[arg(long)]
+    pub package: Vec<PackageName>,
+
     /// Do not de-duplicate repeated dependencies.
     /// Usually, when a package has already displayed its dependencies,
     /// further occurrences will not re-display its dependencies,

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

@@ -25,6 +25,7 @@ use crate::printer::Printer;
 pub(crate) fn pip_tree(
     depth: u8,
     prune: Vec<PackageName>,
+    package: Vec<PackageName>,
     no_dedupe: bool,
     invert: bool,
     strict: bool,
@@ -53,6 +54,7 @@ pub(crate) fn pip_tree(
         &site_packages,
         depth.into(),
         prune,
+        package,
         no_dedupe,
         invert,
         environment.interpreter().markers(),
@@ -117,6 +119,8 @@ struct DisplayDependencyGraph<'env> {
     depth: usize,
     /// Prune the given packages from the display of the dependency tree.
     prune: Vec<PackageName>,
+    /// Display only the specified packages.
+    package: Vec<PackageName>,
     /// Whether to de-duplicate the displayed dependencies.
     no_dedupe: bool,
     /// Map from package name to its requirements.
@@ -131,6 +135,7 @@ impl<'env> DisplayDependencyGraph<'env> {
         site_packages: &'env SitePackages,
         depth: usize,
         prune: Vec<PackageName>,
+        package: Vec<PackageName>,
         no_dedupe: bool,
         invert: bool,
         markers: &'env MarkerEnvironment,
@@ -158,6 +163,7 @@ impl<'env> DisplayDependencyGraph<'env> {
             site_packages,
             depth,
             prune,
+            package,
             no_dedupe,
             requirements,
         })
@@ -260,17 +266,28 @@ impl<'env> DisplayDependencyGraph<'env> {
         let mut path: Vec<&PackageName> = Vec::new();
         let mut lines: Vec<String> = Vec::new();
 
-        // 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 !children.contains(site_package.name()) {
-                path.clear();
-                lines.extend(self.visit(site_package, &mut visited, &mut path)?);
+        if self.package.is_empty() {
+            // 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 !children.contains(site_package.name()) {
+                    path.clear();
+                    lines.extend(self.visit(site_package, &mut visited, &mut path)?);
+                }
+            }
+        } else {
+            for (index, package) in self.package.iter().enumerate() {
+                if index != 0 {
+                    lines.push(String::new());
+                }
+                for installed_dist in self.site_packages.get_packages(package) {
+                    path.clear();
+                    lines.extend(self.visit(installed_dist, &mut visited, &mut path)?);
+                }
             }
         }
-
         Ok(lines)
     }
 }

crates/uv/src/main.rs

@@ -550,6 +550,7 @@ async fn run() -> Result<ExitStatus> {
             commands::pip_tree(
                 args.depth,
                 args.prune,
+                args.package,
                 args.no_dedupe,
                 args.invert,
                 args.shared.strict,

crates/uv/src/settings.rs

@@ -1066,6 +1066,7 @@ impl PipShowSettings {
 pub(crate) struct PipTreeSettings {
     pub(crate) depth: u8,
     pub(crate) prune: Vec<PackageName>,
+    pub(crate) package: Vec<PackageName>,
     pub(crate) no_dedupe: bool,
     pub(crate) invert: bool,
     // CLI-only settings.
@@ -1078,6 +1079,7 @@ impl PipTreeSettings {
         let PipTreeArgs {
             depth,
             prune,
+            package,
             no_dedupe,
             invert,
             strict,
@@ -1091,6 +1093,7 @@ impl PipTreeSettings {
         Self {
             depth,
             prune,
+            package,
             no_dedupe,
             invert,
             // Shared settings.

crates/uv/tests/pip_tree.rs

@@ -1354,3 +1354,153 @@ fn with_editable() {
     "###
     );
 }
+
+#[test]
+#[cfg(target_os = "macos")]
+fn package_flag_complex() {
+    let context = TestContext::new("3.12");
+
+    let requirements_txt = context.temp_dir.child("requirements.txt");
+    requirements_txt.write_str("packse").unwrap();
+
+    uv_snapshot!(context
+        .pip_install()
+        .arg("-r")
+        .arg("requirements.txt")
+        .arg("--strict"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 32 packages in [TIME]
+    Prepared 32 packages in [TIME]
+    Installed 32 packages in [TIME]
+     + certifi==2024.2.2
+     + charset-normalizer==3.3.2
+     + chevron-blue==0.2.1
+     + docutils==0.20.1
+     + hatchling==1.22.4
+     + idna==3.6
+     + importlib-metadata==7.1.0
+     + jaraco-classes==3.3.1
+     + jaraco-context==4.3.0
+     + jaraco-functools==4.0.0
+     + keyring==25.0.0
+     + markdown-it-py==3.0.0
+     + mdurl==0.1.2
+     + more-itertools==10.2.0
+     + msgspec==0.18.6
+     + nh3==0.2.15
+     + packaging==24.0
+     + packse==0.3.12
+     + pathspec==0.12.1
+     + pkginfo==1.10.0
+     + pluggy==1.4.0
+     + pygments==2.17.2
+     + readme-renderer==43.0
+     + requests==2.31.0
+     + requests-toolbelt==1.0.0
+     + rfc3986==2.0.0
+     + rich==13.7.1
+     + setuptools==69.2.0
+     + trove-classifiers==2024.3.3
+     + twine==4.0.2
+     + urllib3==2.2.1
+     + zipp==3.18.1
+    "###
+    );
+
+    uv_snapshot!(
+        context.filters(),
+        context.pip_tree()
+        .arg("--package")
+        .arg("hatchling")
+        .arg("--package")
+        .arg("keyring"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    hatchling v1.22.4
+    ├── packaging v24.0
+    ├── pathspec v0.12.1
+    ├── pluggy v1.4.0
+    └── trove-classifiers v2024.3.3
+
+    keyring v25.0.0
+    ├── jaraco-classes v3.3.1
+    │   └── more-itertools v10.2.0
+    ├── jaraco-functools v4.0.0
+    │   └── more-itertools v10.2.0
+    └── jaraco-context v4.3.0
+
+    ----- stderr -----
+    "###
+    );
+}
+
+#[test]
+fn package_flag() {
+    let context = TestContext::new("3.12");
+
+    let requirements_txt = context.temp_dir.child("requirements.txt");
+    requirements_txt
+        .write_str("scikit-learn==1.4.1.post1")
+        .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]
+     + joblib==1.3.2
+     + numpy==1.26.4
+     + scikit-learn==1.4.1.post1
+     + scipy==1.12.0
+     + threadpoolctl==3.4.0
+    "###
+    );
+
+    uv_snapshot!(
+        context.filters(),
+        context.pip_tree()
+        .arg("--package")
+        .arg("numpy"),
+        @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    numpy v1.26.4
+
+    ----- stderr -----
+    "###
+    );
+
+    uv_snapshot!(
+        context.filters(),
+        context.pip_tree()
+        .arg("--package")
+        .arg("scipy")
+        .arg("--package")
+        .arg("joblib"),
+        @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    scipy v1.12.0
+    └── numpy v1.26.4
+
+    joblib v1.3.2
+
+    ----- stderr -----
+    "###
+    );
+}