Respect extras when limiting JSON tree depth

crates/uv-resolver/src/lock/tree.rs

@@ -655,26 +655,43 @@ impl<'env> TreeDisplay<'env> {
         serde_json::to_string_pretty(&JsonGraph::from(self))
     }
 
-    /// Return the shortest distance from any displayed root to every package within the requested
+    /// Return the packages and edges reachable from the displayed roots within the requested
     /// depth.
-    fn json_distances(&self) -> FxHashMap<NodeIndex, usize> {
+    fn json_traversal(&self) -> JsonTraversal {
         let mut distances = FxHashMap::default();
         let mut queue = VecDeque::new();
+        let mut nodes = FxHashSet::default();
+        let mut edges = FxHashSet::default();
 
         for root in &self.roots {
             match self.graph[*root] {
                 Node::Root => {
                     for edge in self.graph.edges_directed(*root, Direction::Outgoing) {
-                        if matches!(self.graph[edge.target()], Node::Package(_))
-                            && distances.insert(edge.target(), 0).is_none()
-                        {
-                            queue.push_back(edge.target());
+                        let Node::Package(package_id) = self.graph[edge.target()] else {
+                            continue;
+                        };
+                        let state = JsonTraversalNode {
+                            index: edge.target(),
+                            expanded_extras: self.expanded_extras(
+                                self.lock.find_by_id(package_id),
+                                Some(edge.weight()),
+                            ),
+                        };
+                        nodes.insert(state.index);
+                        if distances.insert(state.clone(), 0).is_none() {
+                            queue.push_back(state);
                         }
                     }
                 }
-                Node::Package(_) => {
-                    if distances.insert(*root, 0).is_none() {
-                        queue.push_back(*root);
+                Node::Package(package_id) => {
+                    let state = JsonTraversalNode {
+                        index: *root,
+                        expanded_extras: self
+                            .expanded_extras(self.lock.find_by_id(package_id), None),
+                    };
+                    nodes.insert(state.index);
+                    if distances.insert(state.clone(), 0).is_none() {
+                        queue.push_back(state);
                     }
                 }
             }
@@ -686,22 +703,48 @@ impl<'env> TreeDisplay<'env> {
                 continue;
             }
 
-            for edge in self.graph.edges_directed(source, Direction::Outgoing) {
-                let target = edge.target();
-                if !matches!(self.graph[target], Node::Package(_))
-                    || distances.contains_key(&target)
+            for edge in self.graph.edges_directed(source.index, Direction::Outgoing) {
+                if !self.invert
+                    && let Edge::Optional(required_extra, _) = edge.weight()
+                    && !source.expanded_extras.contains(required_extra)
                 {
                     continue;
                 }
-                distances.insert(target, distance + 1);
-                queue.push_back(target);
+
+                let target = edge.target();
+                let Node::Package(package_id) = self.graph[target] else {
+                    continue;
+                };
+                let state = JsonTraversalNode {
+                    index: target,
+                    expanded_extras: self
+                        .expanded_extras(self.lock.find_by_id(package_id), Some(edge.weight())),
+                };
+                nodes.insert(state.index);
+                edges.insert(edge.id());
+                if !distances.contains_key(&state) {
+                    distances.insert(state.clone(), distance + 1);
+                    queue.push_back(state);
+                }
             }
         }
 
-        distances
+        JsonTraversal { nodes, edges }
     }
 }
 
+#[derive(Debug)]
+struct JsonTraversal {
+    nodes: FxHashSet<NodeIndex>,
+    edges: FxHashSet<EdgeIndex>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct JsonTraversalNode<'env> {
+    index: NodeIndex,
+    expanded_extras: BTreeSet<&'env ExtraName>,
+}
+
 #[derive(Debug, Serialize)]
 struct JsonGraph<'env> {
     directed: bool,
@@ -713,9 +756,9 @@ struct JsonGraph<'env> {
 
 impl<'env> From<&TreeDisplay<'env>> for JsonGraph<'env> {
     fn from(tree: &TreeDisplay<'env>) -> Self {
-        let distances = tree.json_distances();
+        let traversal = tree.json_traversal();
 
-        let mut package_nodes = distances.keys().copied().collect::<Vec<_>>();
+        let mut package_nodes = traversal.nodes.iter().copied().collect::<Vec<_>>();
         package_nodes.sort_by_key(|index| &tree.graph[*index]);
 
         let node_ids = package_nodes
@@ -770,13 +813,7 @@ impl<'env> From<&TreeDisplay<'env>> for JsonGraph<'env> {
         let mut graph_edges = tree
             .graph
             .edge_references()
-            .filter(|edge| {
-                distances
-                    .get(&edge.source())
-                    .is_some_and(|distance| *distance < tree.depth)
-                    && node_ids.contains_key(&edge.source())
-                    && node_ids.contains_key(&edge.target())
-            })
+            .filter(|edge| traversal.edges.contains(&edge.id()))
             .collect::<Vec<_>>();
         graph_edges.sort_by(|left, right| {
             (

crates/uv/tests/project/tree.rs

@@ -1,8 +1,8 @@
-use anyhow::Result;
+use anyhow::{Context, Result};
 use assert_cmd::assert::OutputAssertExt;
 use assert_fs::prelude::*;
 use indoc::{formatdoc, indoc};
-use insta::assert_snapshot;
+use insta::{assert_json_snapshot, assert_snapshot};
 use url::Url;
 
 use uv_static::EnvVars;
@@ -434,6 +434,112 @@ fn json_output_virtual_root() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn json_output_depth_with_extra_context() -> Result<()> {
+    let context = uv_test::test_context!("3.12");
+
+    context.temp_dir.child("pyproject.toml").write_str(
+        r#"
+        [project]
+        name = "project"
+        version = "0.1.0"
+        requires-python = ">=3.12"
+        dependencies = ["package-a", "package-c"]
+
+        [tool.uv.sources]
+        package-a = { path = "packages/package-a" }
+        package-c = { path = "packages/package-c" }
+        "#,
+    )?;
+
+    let package_a = context.temp_dir.child("packages/package-a");
+    package_a.create_dir_all()?;
+    package_a.child("pyproject.toml").write_str(
+        r#"
+        [project]
+        name = "package-a"
+        version = "1.0.0"
+        requires-python = ">=3.12"
+
+        [project.optional-dependencies]
+        feature = ["package-b"]
+
+        [tool.uv.sources]
+        package-b = { path = "../package-b" }
+        "#,
+    )?;
+
+    let package_c = context.temp_dir.child("packages/package-c");
+    package_c.create_dir_all()?;
+    package_c.child("pyproject.toml").write_str(
+        r#"
+        [project]
+        name = "package-c"
+        version = "1.0.0"
+        requires-python = ">=3.12"
+        dependencies = ["package-a[feature]"]
+
+        [tool.uv.sources]
+        package-a = { path = "../package-a" }
+        "#,
+    )?;
+
+    let package_b = context.temp_dir.child("packages/package-b");
+    package_b.create_dir_all()?;
+    package_b.child("pyproject.toml").write_str(
+        r#"
+        [project]
+        name = "package-b"
+        version = "1.0.0"
+        requires-python = ">=3.12"
+        "#,
+    )?;
+
+    let package_names = |depth: u8| -> Result<Vec<String>> {
+        let output = context
+            .tree()
+            .arg("--preview-features")
+            .arg("json-output")
+            .arg("--format")
+            .arg("json")
+            .arg("--universal")
+            .arg("--depth")
+            .arg(depth.to_string())
+            .output()?;
+        output.clone().assert().success();
+        let report: serde_json::Value = serde_json::from_slice(&output.stdout)?;
+        report["nodes"]
+            .as_array()
+            .context("dependency graph nodes should be an array")?
+            .iter()
+            .map(|node| {
+                node["name"]
+                    .as_str()
+                    .context("dependency graph node should have a name")
+                    .map(ToOwned::to_owned)
+            })
+            .collect()
+    };
+
+    assert_json_snapshot!(package_names(2)?, @r#"
+    [
+      "package-a",
+      "package-c",
+      "project"
+    ]
+    "#);
+    assert_json_snapshot!(package_names(3)?, @r#"
+    [
+      "package-a",
+      "package-b",
+      "package-c",
+      "project"
+    ]
+    "#);
+
+    Ok(())
+}
+
 #[test]
 fn nested_platform_dependencies() -> Result<()> {
     let context = uv_test::test_context!("3.12");