implement --depth, --prune for pip tree (#4440)

crates/uv-cli/src/lib.rs

@@ -1404,6 +1404,13 @@ pub struct PipShowArgs {
 #[derive(Args)]
 #[allow(clippy::struct_excessive_bools)]
 pub struct PipTreeArgs {
+    /// Maximum display depth of the dependency tree
+    #[arg(long, short, default_value_t = 255)]
+    pub depth: u8,
+
+    /// Prune the given package from the display of the dependency tree.
+    #[arg(long)]
+    pub prune: 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

@@ -19,7 +19,10 @@ use std::collections::{HashMap, HashSet};
 use pypi_types::VerbatimParsedUrl;
 
 /// Display the installed packages in the current environment as a dependency tree.
+#[allow(clippy::too_many_arguments)]
 pub(crate) fn pip_tree(
+    depth: u8,
+    prune: Vec<PackageName>,
     no_dedupe: bool,
     strict: bool,
     python: Option<&str>,
@@ -44,7 +47,7 @@ pub(crate) fn pip_tree(
     // Build the installed index.
     let site_packages = SitePackages::from_environment(&environment)?;
 
-    let rendered_tree = DisplayDependencyGraph::new(&site_packages, no_dedupe)
+    let rendered_tree = DisplayDependencyGraph::new(&site_packages, depth.into(), prune, no_dedupe)
         .render()
         .join("\n");
     writeln!(printer.stdout(), "{rendered_tree}").unwrap();
@@ -105,13 +108,24 @@ struct DisplayDependencyGraph<'a> {
     // dependency graph.
     required_packages: HashSet<PackageName>,
 
+    // Maximum display depth of the dependency tree
+    depth: usize,
+
+    // Prune the given package from the display of the dependency tree.
+    prune: Vec<PackageName>,
+
     // Whether to de-duplicate the displayed dependencies.
     no_dedupe: bool,
 }
 
 impl<'a> DisplayDependencyGraph<'a> {
     /// Create a new [`DisplayDependencyGraph`] for the set of installed distributions.
-    fn new(site_packages: &'a SitePackages, no_dedupe: bool) -> DisplayDependencyGraph<'a> {
+    fn new(
+        site_packages: &'a SitePackages,
+        depth: usize,
+        prune: Vec<PackageName>,
+        no_dedupe: bool,
+    ) -> DisplayDependencyGraph<'a> {
         let mut dist_by_package_name = HashMap::new();
         let mut required_packages = HashSet::new();
         for site_package in site_packages.iter() {
@@ -127,6 +141,8 @@ impl<'a> DisplayDependencyGraph<'a> {
             site_packages,
             dist_by_package_name,
             required_packages,
+            depth,
+            prune,
             no_dedupe,
         }
     }
@@ -138,6 +154,16 @@ impl<'a> DisplayDependencyGraph<'a> {
         visited: &mut HashSet<String>,
         path: &mut Vec<String>,
     ) -> Vec<String> {
+        // Short-circuit if the current path is longer than the provided depth.
+        if path.len() > self.depth {
+            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());

crates/uv/src/main.rs

@@ -548,6 +548,8 @@ async fn run() -> Result<ExitStatus> {
             let cache = cache.init()?;
 
             commands::pip_tree(
+                args.depth,
+                args.prune,
                 args.no_dedupe,
                 args.shared.strict,
                 args.shared.python.as_deref(),

crates/uv/src/settings.rs

@@ -1009,6 +1009,8 @@ impl PipShowSettings {
 #[allow(clippy::struct_excessive_bools)]
 #[derive(Debug, Clone)]
 pub(crate) struct PipTreeSettings {
+    pub(crate) depth: u8,
+    pub(crate) prune: Vec<PackageName>,
     pub(crate) no_dedupe: bool,
     // CLI-only settings.
     pub(crate) shared: PipSettings,
@@ -1018,6 +1020,8 @@ impl PipTreeSettings {
     /// Resolve the [`PipTreeSettings`] from the CLI and workspace configuration.
     pub(crate) fn resolve(args: PipTreeArgs, filesystem: Option<FilesystemOptions>) -> Self {
         let PipTreeArgs {
+            depth,
+            prune,
             no_dedupe,
             strict,
             no_strict,
@@ -1027,6 +1031,8 @@ impl PipTreeSettings {
         } = args;
 
         Self {
+            depth,
+            prune,
             no_dedupe,
             // Shared settings.
             shared: PipSettings::combine(

crates/uv/tests/pip_tree.rs

@@ -122,6 +122,206 @@ fn nested_dependencies() {
     );
 }
 
+#[test]
+fn depth() {
+    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!(install_command(&context)
+        .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(), Command::new(get_bin())
+        .arg("pip")
+        .arg("tree")
+        .arg("--cache-dir")
+        .arg(context.cache_dir.path())
+        .arg("--depth")
+        .arg("0")
+        .env("VIRTUAL_ENV", context.venv.as_os_str())
+        .env("UV_NO_WRAP", "1")
+        .current_dir(&context.temp_dir), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    scikit-learn v1.4.1.post1
+
+    ----- stderr -----
+
+    "###
+    );
+
+    uv_snapshot!(context.filters(), Command::new(get_bin())
+        .arg("pip")
+        .arg("tree")
+        .arg("--cache-dir")
+        .arg(context.cache_dir.path())
+        .arg("--depth")
+        .arg("1")
+        .env("VIRTUAL_ENV", context.venv.as_os_str())
+        .env("UV_NO_WRAP", "1")
+        .current_dir(&context.temp_dir), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    scikit-learn v1.4.1.post1
+    ├── numpy v1.26.4
+    ├── scipy v1.12.0
+    ├── joblib v1.3.2
+    └── threadpoolctl v3.4.0
+
+    ----- stderr -----
+
+    "###
+    );
+
+    uv_snapshot!(context.filters(), Command::new(get_bin())
+        .arg("pip")
+        .arg("tree")
+        .arg("--cache-dir")
+        .arg(context.cache_dir.path())
+        .arg("--depth")
+        .arg("2")
+        .env("VIRTUAL_ENV", context.venv.as_os_str())
+        .env("UV_NO_WRAP", "1")
+        .current_dir(&context.temp_dir), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    scikit-learn v1.4.1.post1
+    ├── numpy v1.26.4
+    ├── scipy v1.12.0
+    │   └── numpy v1.26.4 (*)
+    ├── joblib v1.3.2
+    └── threadpoolctl v3.4.0
+    (*) Package tree already displayed
+
+    ----- stderr -----
+
+    "###
+    );
+}
+
+#[test]
+fn prune() {
+    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!(install_command(&context)
+        .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(), Command::new(get_bin())
+        .arg("pip")
+        .arg("tree")
+        .arg("--cache-dir")
+        .arg(context.cache_dir.path())
+        .arg("--prune")
+        .arg("numpy")
+        .env("VIRTUAL_ENV", context.venv.as_os_str())
+        .env("UV_NO_WRAP", "1")
+        .current_dir(&context.temp_dir), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    scikit-learn v1.4.1.post1
+    ├── scipy v1.12.0
+    ├── joblib v1.3.2
+    └── threadpoolctl v3.4.0
+
+    ----- stderr -----
+
+    "###
+    );
+
+    uv_snapshot!(context.filters(), Command::new(get_bin())
+        .arg("pip")
+        .arg("tree")
+        .arg("--cache-dir")
+        .arg(context.cache_dir.path())
+        .arg("--prune")
+        .arg("numpy")
+        .arg("--prune")
+        .arg("joblib")
+        .env("VIRTUAL_ENV", context.venv.as_os_str())
+        .env("UV_NO_WRAP", "1")
+        .current_dir(&context.temp_dir), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    scikit-learn v1.4.1.post1
+    ├── scipy v1.12.0
+    └── threadpoolctl v3.4.0
+
+    ----- stderr -----
+
+    "###
+    );
+
+    uv_snapshot!(context.filters(), Command::new(get_bin())
+        .arg("pip")
+        .arg("tree")
+        .arg("--cache-dir")
+        .arg(context.cache_dir.path())
+        .arg("--prune")
+        .arg("scipy")
+        .env("VIRTUAL_ENV", context.venv.as_os_str())
+        .env("UV_NO_WRAP", "1")
+        .current_dir(&context.temp_dir), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    scikit-learn v1.4.1.post1
+    ├── numpy v1.26.4
+    ├── joblib v1.3.2
+    └── threadpoolctl v3.4.0
+
+    ----- stderr -----
+
+    "###
+    );
+}
+
 #[test]
 #[cfg(target_os = "macos")]
 fn nested_dependencies_more_complex() {
@@ -225,6 +425,112 @@ fn nested_dependencies_more_complex() {
     );
 }
 
+#[test]
+#[cfg(target_os = "macos")]
+fn prune_big_tree() {
+    let context = TestContext::new("3.12");
+
+    let requirements_txt = context.temp_dir.child("requirements.txt");
+    requirements_txt.write_str("packse").unwrap();
+
+    uv_snapshot!(install_command(&context)
+        .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(), Command::new(get_bin())
+        .arg("pip")
+        .arg("tree")
+        .arg("--cache-dir")
+        .arg(context.cache_dir.path())
+        .arg("--prune")
+        .arg("hatchling")
+        .env("VIRTUAL_ENV", context.venv.as_os_str())
+        .env("UV_NO_WRAP", "1")
+        .current_dir(&context.temp_dir), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    packse v0.3.12
+    ├── chevron-blue v0.2.1
+    ├── msgspec v0.18.6
+    ├── setuptools v69.2.0
+    └── twine v4.0.2
+        ├── pkginfo v1.10.0
+        ├── readme-renderer v43.0
+        │   ├── nh3 v0.2.15
+        │   ├── docutils v0.20.1
+        │   └── pygments v2.17.2
+        ├── requests v2.31.0
+        │   ├── charset-normalizer v3.3.2
+        │   ├── idna v3.6
+        │   ├── urllib3 v2.2.1
+        │   └── certifi v2024.2.2
+        ├── requests-toolbelt v1.0.0
+        │   └── requests v2.31.0 (*)
+        ├── urllib3 v2.2.1 (*)
+        ├── importlib-metadata v7.1.0
+        │   └── zipp v3.18.1
+        ├── 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
+        ├── rfc3986 v2.0.0
+        └── rich v13.7.1
+            ├── markdown-it-py v3.0.0
+            │   └── mdurl v0.1.2
+            └── pygments v2.17.2 (*)
+    (*) Package tree already displayed
+
+    ----- stderr -----
+    "###
+    );
+}
+
 // Ensure `pip tree` behaves correctly with a package that has a cyclic dependency.
 // package `uv-cyclic-dependencies-a` and `uv-cyclic-dependencies-b` depend on each other,
 // which creates a dependency cycle.