Support --depth for graphviz output format

src/pipdeptree/_cli.py

@@ -121,7 +121,7 @@ def build_parser() -> ArgumentParser:
         "--depth",
         type=lambda x: int(x) if x.isdigit() and (int(x) >= 0) else parser.error("Depth must be a number that is >= 0"),
         default=float("inf"),
-        help="limit the depth of the tree (text and freeze render only)",
+        help="limit the depth of the tree (text, freeze, and graphviz render only)",
         metavar="D",
     )
     render.add_argument(

src/pipdeptree/_render/__init__.py

@@ -25,7 +25,9 @@ def render(options: Options, tree: PackageDAG) -> None:
     elif output_format == "freeze":
         render_freeze(tree, max_depth=options.depth, list_all=options.all)
     elif output_format.startswith("graphviz-"):
-        render_graphviz(tree, output_format=output_format[len("graphviz-") :], reverse=options.reverse)
+        render_graphviz(
+            tree, output_format=output_format[len("graphviz-") :], reverse=options.reverse, max_depth=options.depth
+        )
     else:
         render_text(
             tree,

src/pipdeptree/_render/graphviz.py

@@ -1,26 +1,125 @@
 from __future__ import annotations
 
+import math
 import os
 import sys
+from collections import deque
 from typing import TYPE_CHECKING
 
 from pipdeptree._models import DistPackage, ReqPackage
 
 if TYPE_CHECKING:
+    from graphviz import Digraph
+
     from pipdeptree._models import PackageDAG
 
 
-def dump_graphviz(  # noqa: C901
+def _get_all_dep_keys(tree: PackageDAG) -> set[str]:
+    """Return the set of keys that appear as dependencies of some package."""
+    dep_keys: set[str] = set()
+    for deps in tree.values():
+        dep_keys.update(dep.key for dep in deps)
+    return dep_keys
+
+
+def _build_reverse_graph(tree: PackageDAG, graph: Digraph, max_depth: float) -> None:  # noqa: C901, PLR0912
+    """Build graphviz nodes and edges for a reversed dependency tree."""
+    if max_depth < math.inf:
+        parent_keys = _get_all_dep_keys(tree)
+        root_keys = {dep_rev.key for dep_rev in tree if dep_rev.key not in parent_keys}
+        visited: dict[str, int] = {}
+        queue: deque[tuple[str, int]] = deque((k, 0) for k in root_keys)
+        while queue:
+            key, depth = queue.popleft()
+            if key in visited:
+                continue
+            visited[key] = depth
+            if depth < max_depth:
+                for parent in tree.get_children(key):
+                    if parent.key not in visited:
+                        queue.append((parent.key, depth + 1))
+        for dep_rev, parents in tree.items():
+            if dep_rev.key not in visited:
+                continue
+            assert isinstance(dep_rev, ReqPackage)
+            dep_label = f"{dep_rev.project_name}\\n{dep_rev.installed_version}"
+            graph.node(dep_rev.key, label=dep_label)
+            if visited[dep_rev.key] < max_depth:
+                for parent in parents:
+                    assert isinstance(parent, DistPackage)
+                    if parent.key in visited:
+                        edge_label = (parent.req.version_spec if parent.req is not None else None) or "any"
+                        graph.edge(dep_rev.key, parent.key, label=edge_label)
+    else:
+        for dep_rev, parents in tree.items():
+            assert isinstance(dep_rev, ReqPackage)
+            dep_label = f"{dep_rev.project_name}\\n{dep_rev.installed_version}"
+            graph.node(dep_rev.key, label=dep_label)
+            for parent in parents:
+                assert isinstance(parent, DistPackage)
+                edge_label = (parent.req.version_spec if parent.req is not None else None) or "any"
+                graph.edge(dep_rev.key, parent.key, label=edge_label)
+
+
+def _build_forward_graph(tree: PackageDAG, graph: Digraph, max_depth: float) -> None:  # noqa: C901, PLR0912
+    """Build graphviz nodes and edges for a forward dependency tree."""
+    if max_depth < math.inf:  # noqa: PLR1702
+        dep_keys = _get_all_dep_keys(tree)
+        root_keys = {pkg.key for pkg in tree if pkg.key not in dep_keys}
+        visited: dict[str, int] = {}
+        queue: deque[tuple[str, int]] = deque((k, 0) for k in root_keys)
+        while queue:
+            key, depth = queue.popleft()
+            if key in visited:
+                continue
+            visited[key] = depth
+            if depth < max_depth:
+                children = tree.get_children(key)
+                for dep in children:
+                    if dep.key not in visited:
+                        queue.append((dep.key, depth + 1))
+        for pkg, deps in tree.items():
+            if pkg.key not in visited:
+                continue
+            pkg_label = f"{pkg.project_name}\\n{pkg.version}"
+            graph.node(pkg.key, label=pkg_label)
+            if visited[pkg.key] < max_depth:
+                for dep in deps:
+                    if dep.key in visited:
+                        edge_label = dep.version_spec or "any"
+                        if dep.is_missing:
+                            dep_label = f"{dep.project_name}\\n(missing)"
+                            graph.node(dep.key, label=dep_label, style="dashed")
+                            graph.edge(pkg.key, dep.key, style="dashed")
+                        else:
+                            graph.edge(pkg.key, dep.key, label=edge_label)
+    else:
+        for pkg, deps in tree.items():
+            pkg_label = f"{pkg.project_name}\\n{pkg.version}"
+            graph.node(pkg.key, label=pkg_label)
+            for dep in deps:
+                edge_label = dep.version_spec or "any"
+                if dep.is_missing:
+                    dep_label = f"{dep.project_name}\\n(missing)"
+                    graph.node(dep.key, label=dep_label, style="dashed")
+                    graph.edge(pkg.key, dep.key, style="dashed")
+                else:
+                    graph.edge(pkg.key, dep.key, label=edge_label)
+
+
+def dump_graphviz(
     tree: PackageDAG,
     output_format: str = "dot",
     is_reverse: bool = False,  # noqa: FBT001, FBT002
+    max_depth: float = math.inf,
 ) -> str | bytes:
     """
     Output dependency graph as one of the supported GraphViz output formats.
 
     :param dict tree: dependency graph
     :param string output_format: output format
     :param bool is_reverse: reverse or not
+    :param float max_depth: maximum depth of the dependency tree to include
     :returns: representation of tree in the specified output format
     :rtype: str or binary representation depending on the output format
     """
@@ -45,27 +144,9 @@ def dump_graphviz(  # noqa: C901
     graph = Digraph(format=output_format)
 
     if is_reverse:
-        for dep_rev, parents in tree.items():
-            assert isinstance(dep_rev, ReqPackage)
-            dep_label = f"{dep_rev.project_name}\\n{dep_rev.installed_version}"
-            graph.node(dep_rev.key, label=dep_label)
-            for parent in parents:
-                # req reference of the dep associated with this particular parent package
-                assert isinstance(parent, DistPackage)
-                edge_label = (parent.req.version_spec if parent.req is not None else None) or "any"
-                graph.edge(dep_rev.key, parent.key, label=edge_label)
+        _build_reverse_graph(tree, graph, max_depth)
     else:
-        for pkg, deps in tree.items():
-            pkg_label = f"{pkg.project_name}\\n{pkg.version}"
-            graph.node(pkg.key, label=pkg_label)
-            for dep in deps:
-                edge_label = dep.version_spec or "any"
-                if dep.is_missing:
-                    dep_label = f"{dep.project_name}\\n(missing)"
-                    graph.node(dep.key, label=dep_label, style="dashed")
-                    graph.edge(pkg.key, dep.key, style="dashed")
-                else:
-                    graph.edge(pkg.key, dep.key, label=edge_label)
+        _build_forward_graph(tree, graph, max_depth)
 
     # Allow output of dot format, even if GraphViz isn't installed.
     if output_format == "dot":
@@ -97,8 +178,8 @@ def print_graphviz(dump_output: str | bytes) -> None:
             bytestream.write(dump_output)
 
 
-def render_graphviz(tree: PackageDAG, *, output_format: str, reverse: bool) -> None:
-    output = dump_graphviz(tree, output_format=output_format, is_reverse=reverse)
+def render_graphviz(tree: PackageDAG, *, output_format: str, reverse: bool, max_depth: float = math.inf) -> None:
+    output = dump_graphviz(tree, output_format=output_format, is_reverse=reverse, max_depth=max_depth)
     print_graphviz(output)
 
 

tests/render/test_graphviz.py

@@ -68,3 +68,61 @@ def test_render_svg(capsys: pytest.CaptureFixture[str], example_dag: PackageDAG)
     assert out.startswith("<?xml")
     assert "<svg" in out
     assert out.strip().endswith("</svg>")
+
+
+def test_render_dot_with_depth(example_dag: PackageDAG) -> None:
+    output = dump_graphviz(example_dag, output_format="dot", max_depth=1)
+    assert isinstance(output, str)
+    # Roots are a and g (depth 0), their direct deps are at depth 1
+    # Deeper deps (d) should not appear
+    assert "a ->" in output
+    assert "g ->" in output
+    assert 'd [label="d' not in output  # d is only reachable at depth 2+
+    # Direct deps of roots should appear
+    assert 'b [label="b' in output
+    assert 'c [label="c' in output
+    assert 'e [label="e' in output
+    assert 'f [label="f' in output
+
+
+def test_render_dot_reverse_with_depth(example_dag: PackageDAG) -> None:
+    reversed_dag = example_dag.reverse()
+    output = dump_graphviz(reversed_dag, output_format="dot", is_reverse=True, max_depth=1)
+    assert isinstance(output, str)
+    # In reverse, root is e (only leaf in forward tree that's not a parent in reverse)
+    # At depth 0: e, at depth 1: c, d, g (e's parents)
+    assert 'e [label="e' in output
+    assert 'c [label="c' in output
+    assert 'd [label="d' in output
+    assert 'g [label="g' in output
+    # Deeper nodes should not appear
+    assert 'a [label="a' not in output
+    assert 'b [label="b' not in output
+    assert 'f [label="f' not in output
+    # Edges from e to its parents should exist
+    assert "e -> c" in output
+    assert "e -> d" in output
+    assert "e -> g" in output
+
+
+def test_render_dot_reverse_infinite_depth(example_dag: PackageDAG) -> None:
+    reversed_dag = example_dag.reverse()
+    output = dump_graphviz(reversed_dag, output_format="dot", is_reverse=True)
+    assert isinstance(output, str)
+    # All nodes should appear with no depth limit
+    for node in ("a", "b", "c", "d", "e", "f", "g"):
+        assert f'{node} [label="' in output
+    # Edges should include reverse relationships
+    assert "e -> c" in output
+    assert "e -> d" in output
+    assert "b -> a" in output
+    assert "f -> g" in output
+
+
+def test_render_dot_with_depth_zero(example_dag: PackageDAG) -> None:
+    output = dump_graphviz(example_dag, output_format="dot", max_depth=0)
+    assert isinstance(output, str)
+    # Depth 0 means only root nodes, no edges
+    assert "a [label=" in output
+    assert "g [label=" in output
+    assert "->" not in output

tests/render/test_render.py

@@ -37,7 +37,7 @@ def test_mermaid_routing(option: list[str], mocker: MockerFixture) -> None:
 def test_grahpviz_routing(option: list[str], mocker: MockerFixture) -> None:
     render = mocker.patch("pipdeptree._render.render_graphviz")
     main(option)
-    render.assert_called_once_with(ANY, output_format="dot", reverse=False)
+    render.assert_called_once_with(ANY, output_format="dot", reverse=False, max_depth=inf)
 
 
 @pytest.mark.parametrize("option", [[], ["--output", "text"]])