feat: Implement --depth workspace for cargo tree command

Author: ShoyuVanilla

src/bin/cargo/commands/tree.rs

@@ -2,7 +2,7 @@ use crate::cli;
 use crate::command_prelude::*;
 use anyhow::{bail, format_err};
 use cargo::core::dependency::DepKind;
-use cargo::ops::tree::{self, EdgeKind};
+use cargo::ops::tree::{self, DisplayDepth, EdgeKind};
 use cargo::ops::Packages;
 use cargo::util::print_available_packages;
 use cargo::util::CargoResult;
@@ -42,7 +42,14 @@ pub fn cli() -> Command {
             "SPEC",
             "Prune the given package from the display of the dependency tree",
         ))
-        .arg(opt("depth", "Maximum display depth of the dependency tree").value_name("DEPTH"))
+        .arg(
+            opt(
+                "depth",
+                "Maximum display depth of the dependency tree, \
+                or --depth=worktree to display dependencies from workspace members only",
+            )
+            .value_name("DEPTH"),
+        )
         .arg(flag("no-indent", "Deprecated, use --prefix=none instead").hide(true))
         .arg(flag("prefix-depth", "Deprecated, use --prefix=depth instead").hide(true))
         .arg(
@@ -162,6 +169,12 @@ pub fn exec(gctx: &mut GlobalContext, args: &ArgMatches) -> CliResult {
 
     let pkgs_to_prune = args._values_of("prune");
 
+    let display_depth = args
+        ._value_of("depth")
+        .map(|s| s.parse::<DisplayDepth>())
+        .transpose()?
+        .unwrap_or(DisplayDepth::MaxDisplayDepth(u32::MAX));
+
     let packages = args.packages_from_flags()?;
     let mut invert = args
         .get_many::<String>("invert")
@@ -222,7 +235,7 @@ subtree of the package given to -p.\n\
         duplicates: args.flag("duplicates"),
         format: args.get_one::<String>("format").cloned().unwrap(),
         graph_features,
-        max_display_depth: args.value_of_u32("depth")?.unwrap_or(u32::MAX),
+        display_depth,
         no_proc_macro,
     };
 

src/cargo/ops/tree/mod.rs

@@ -6,7 +6,7 @@ use crate::core::dependency::DepKind;
 use crate::core::resolver::{features::CliFeatures, ForceAllTargets, HasDevUnits};
 use crate::core::{Package, PackageId, PackageIdSpec, PackageIdSpecQuery, Workspace};
 use crate::ops::{self, Packages};
-use crate::util::{CargoResult, GlobalContext};
+use crate::util::CargoResult;
 use crate::{drop_print, drop_println};
 use anyhow::Context as _;
 use graph::Graph;
@@ -43,8 +43,10 @@ pub struct TreeOptions {
     pub format: String,
     /// Includes features in the tree as separate nodes.
     pub graph_features: bool,
-    /// Maximum display depth of the dependency tree.
-    pub max_display_depth: u32,
+    /// Display depth of the dependency tree.
+    /// If non-negative integer, display dependencies with that amount of max depth.
+    /// If `workspace`, display dependencies from current workspace only.
+    pub display_depth: DisplayDepth,
     /// Excludes proc-macro dependencies.
     pub no_proc_macro: bool,
 }
@@ -86,6 +88,32 @@ impl FromStr for Prefix {
     }
 }
 
+#[derive(Clone, Copy)]
+pub enum DisplayDepth {
+    MaxDisplayDepth(u32),
+    Workspace,
+}
+
+impl FromStr for DisplayDepth {
+    type Err = clap::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "workspace" => Ok(Self::Workspace),
+            s => s.parse().map(Self::MaxDisplayDepth).map_err(|_| {
+                clap::Error::raw(
+                    clap::error::ErrorKind::ValueValidation,
+                    format!(
+                        "supported values for --depth are non-negative integers and `workspace`, \
+                                but `{}` is unknown",
+                        s
+                    ),
+                )
+            }),
+        }
+    }
+}
+
 struct Symbols {
     down: &'static str,
     tee: &'static str,
@@ -203,14 +231,14 @@ pub fn build_and_print(ws: &Workspace<'_>, opts: &TreeOptions) -> CargoResult<()
         try to use option `--target all` first, and then narrow your search scope accordingly.",
         )?;
     } else {
-        print(ws.gctx(), opts, root_indexes, &pkgs_to_prune, &graph)?;
+        print(ws, opts, root_indexes, &pkgs_to_prune, &graph)?;
     }
     Ok(())
 }
 
 /// Prints a tree for each given root.
 fn print(
-    gctx: &GlobalContext,
+    ws: &Workspace<'_>,
     opts: &TreeOptions,
     roots: Vec<usize>,
     pkgs_to_prune: &[PackageIdSpec],
@@ -219,7 +247,7 @@ fn print(
     let format = Pattern::new(&opts.format)
         .with_context(|| format!("tree format `{}` not valid", opts.format))?;
 
-    let symbols = if gctx.shell().out_unicode() {
+    let symbols = if ws.gctx().shell().out_unicode() {
         &UTF8_SYMBOLS
     } else {
         &ASCII_SYMBOLS
@@ -231,7 +259,7 @@ fn print(
 
     for (i, root_index) in roots.into_iter().enumerate() {
         if i != 0 {
-            drop_println!(gctx);
+            drop_println!(ws.gctx());
         }
 
         // A stack of bools used to determine where | symbols should appear
@@ -242,15 +270,15 @@ fn print(
         let mut print_stack = vec![];
 
         print_node(
-            gctx,
+            ws,
             graph,
             root_index,
             &format,
             symbols,
             pkgs_to_prune,
             opts.prefix,
             opts.no_dedupe,
-            opts.max_display_depth,
+            opts.display_depth,
             &mut visited_deps,
             &mut levels_continue,
             &mut print_stack,
@@ -262,36 +290,36 @@ fn print(
 
 /// Prints a package and all of its dependencies.
 fn print_node<'a>(
-    gctx: &GlobalContext,
+    ws: &Workspace<'_>,
     graph: &'a Graph<'_>,
     node_index: usize,
     format: &Pattern,
     symbols: &Symbols,
     pkgs_to_prune: &[PackageIdSpec],
     prefix: Prefix,
     no_dedupe: bool,
-    max_display_depth: u32,
+    display_depth: DisplayDepth,
     visited_deps: &mut HashSet<usize>,
     levels_continue: &mut Vec<bool>,
     print_stack: &mut Vec<usize>,
 ) {
     let new = no_dedupe || visited_deps.insert(node_index);
 
     match prefix {
-        Prefix::Depth => drop_print!(gctx, "{}", levels_continue.len()),
+        Prefix::Depth => drop_print!(ws.gctx(), "{}", levels_continue.len()),
         Prefix::Indent => {
             if let Some((last_continues, rest)) = levels_continue.split_last() {
                 for continues in rest {
                     let c = if *continues { symbols.down } else { " " };
-                    drop_print!(gctx, "{}   ", c);
+                    drop_print!(ws.gctx(), "{}   ", c);
                 }
 
                 let c = if *last_continues {
                     symbols.tee
                 } else {
                     symbols.ell
                 };
-                drop_print!(gctx, "{0}{1}{1} ", c, symbols.right);
+                drop_print!(ws.gctx(), "{0}{1}{1} ", c, symbols.right);
             }
         }
         Prefix::None => {}
@@ -307,7 +335,7 @@ fn print_node<'a>(
     } else {
         " (*)"
     };
-    drop_println!(gctx, "{}{}", format.display(graph, node_index), star);
+    drop_println!(ws.gctx(), "{}{}", format.display(graph, node_index), star);
 
     if !new || in_cycle {
         return;
@@ -321,15 +349,15 @@ fn print_node<'a>(
         EdgeKind::Feature,
     ] {
         print_dependencies(
-            gctx,
+            ws,
             graph,
             node_index,
             format,
             symbols,
             pkgs_to_prune,
             prefix,
             no_dedupe,
-            max_display_depth,
+            display_depth,
             visited_deps,
             levels_continue,
             print_stack,
@@ -341,15 +369,15 @@ fn print_node<'a>(
 
 /// Prints all the dependencies of a package for the given dependency kind.
 fn print_dependencies<'a>(
-    gctx: &GlobalContext,
+    ws: &Workspace<'_>,
     graph: &'a Graph<'_>,
     node_index: usize,
     format: &Pattern,
     symbols: &Symbols,
     pkgs_to_prune: &[PackageIdSpec],
     prefix: Prefix,
     no_dedupe: bool,
-    max_display_depth: u32,
+    display_depth: DisplayDepth,
     visited_deps: &mut HashSet<usize>,
     levels_continue: &mut Vec<bool>,
     print_stack: &mut Vec<usize>,
@@ -371,13 +399,18 @@ fn print_dependencies<'a>(
         if let Some(name) = name {
             for continues in &**levels_continue {
                 let c = if *continues { symbols.down } else { " " };
-                drop_print!(gctx, "{}   ", c);
+                drop_print!(ws.gctx(), "{}   ", c);
             }
 
-            drop_println!(gctx, "{}", name);
+            drop_println!(ws.gctx(), "{}", name);
         }
     }
 
+    let (max_display_depth, filter_non_workspace_member) = match display_depth {
+        DisplayDepth::MaxDisplayDepth(max) => (max, false),
+        DisplayDepth::Workspace => (u32::MAX, true),
+    };
+
     // Current level exceeds maximum display depth. Skip.
     if levels_continue.len() + 1 > max_display_depth as usize {
         return;
@@ -389,6 +422,11 @@ fn print_dependencies<'a>(
             // Filter out packages to prune.
             match graph.node(**dep) {
                 Node::Package { package_id, .. } => {
+                    if filter_non_workspace_member
+                        && !ws.is_member(graph.package_for_id(*package_id))
+                    {
+                        return false;
+                    }
                     !pkgs_to_prune.iter().any(|spec| spec.matches(*package_id))
                 }
                 _ => true,
@@ -399,15 +437,15 @@ fn print_dependencies<'a>(
     while let Some(dependency) = it.next() {
         levels_continue.push(it.peek().is_some());
         print_node(
-            gctx,
+            ws,
             graph,
             *dependency,
             format,
             symbols,
             pkgs_to_prune,
             prefix,
             no_dedupe,
-            max_display_depth,
+            display_depth,
             visited_deps,
             levels_continue,
             print_stack,

tests/testsuite/tree.rs

@@ -1846,6 +1846,61 @@ c v1.0.0
         .run();
 }
 
+#[cargo_test]
+fn depth_workspace() {
+    Package::new("somedep", "1.0.0").publish();
+    Package::new("otherdep", "1.0.0").publish();
+    let p = project()
+        .file(
+            "Cargo.toml",
+            r#"
+            [workspace]
+            members = ["a", "b", "c"]
+            "#,
+        )
+        .file("a/Cargo.toml", &basic_manifest("a", "1.0.0"))
+        .file("a/src/lib.rs", "")
+        .file(
+            "b/Cargo.toml",
+            r#"
+            [package]
+            name = "b"
+            version = "0.1.0"
+
+            [dependencies]
+            c = { path = "../c" }
+            somedep = "1"
+            "#,
+        )
+        .file("b/src/lib.rs", "")
+        .file(
+            "c/Cargo.toml",
+            r#"
+            [package]
+            name = "c"
+            version = "0.1.0"
+
+            [dependencies]
+            somedep = "1"
+            otherdep = "1"
+            "#,
+        )
+        .file("c/src/lib.rs", "")
+        .build();
+
+    p.cargo("tree --depth workspace")
+        .with_stdout_data(str![[r#"
+a v1.0.0 ([ROOT]/foo/a)
+
+b v0.1.0 ([ROOT]/foo/b)
+└── c v0.1.0 ([ROOT]/foo/c)
+
+c v0.1.0 ([ROOT]/foo/c) (*)
+
+"#]])
+        .run();
+}
+
 #[cargo_test]
 fn prune() {
     let p = make_simple_proj();