Add --max-depth option to control diagram complexity (#10077)

Julfried · Pierre-Sassoulas · web-flow · commit 05e83c49d8b3 · 2025-01-24T21:36:17.000+01:00
Co-authored-by: Pierre Sassoulas &lt;pierre.sassoulas@gmail.com&gt;
diff --git a/doc/additional_tools/pyreverse/configuration.rst b/doc/additional_tools/pyreverse/configuration.rst
@@ -58,6 +58,13 @@ Filtering and Scope
 **Default:**  ``PUB_ONLY``
 
 
+--max-depth
+-----------
+*Maximum depth in package/module hierarchy to display. A depth of 0 shows only top-level packages, 1 shows one level of subpackages, etc. If not specified, all packages/modules are shown.*
+
+**Default:**  ``None``
+
+
 --show-ancestors
 ----------------
 *Show <ancestor> generations of ancestor classes not in <projects>.*
diff --git a/doc/whatsnew/fragments/10077.feature b/doc/whatsnew/fragments/10077.feature
@@ -0,0 +1,4 @@
+Add --max-depth option to pyreverse to control diagram complexity. A depth of 0 shows only top-level packages, 1 shows one level of subpackages, etc.
+This helps manage visualization of large codebases by limiting the depth of displayed packages and classes.
+
+Refs #10077
diff --git a/pylint/pyreverse/main.py b/pylint/pyreverse/main.py
@@ -151,6 +151,22 @@
             "help": "Include standard library objects in representation of classes.",
         },
     ),
+    (
+        "max-depth",
+        {
+            "dest": "max_depth",
+            "action": "store",
+            "default": None,
+            "metavar": "<depth>",
+            "type": "int",
+            "group": OPTIONS_GROUPS["FILTERING"],
+            "help": (
+                "Maximum depth in package/module hierarchy to display. A depth of 0 shows only "
+                "top-level packages, 1 shows one level of subpackages, etc. If not specified, "
+                "all packages/modules are shown."
+            ),
+        },
+    ),
     # Display Options
     (
         "module-names",
diff --git a/pylint/pyreverse/writer.py b/pylint/pyreverse/writer.py
@@ -35,6 +35,7 @@ def __init__(self, config: argparse.Namespace) -> None:
         self.printer: Printer  # defined in set_printer
         self.file_name = ""  # defined in set_printer
         self.depth = self.config.max_color_depth
+        self.max_depth = self.config.max_depth
         # default colors are an adaption of the seaborn colorblind palette
         self.available_colors = itertools.cycle(self.config.color_palette)
         self.used_colors: dict[str, str] = {}
@@ -53,6 +54,38 @@ def write(self, diadefs: Iterable[ClassDiagram | PackageDiagram]) -> None:
                 self.write_classes(diagram)
             self.save()
 
+    def should_show_node(self, qualified_name: str, is_class: bool = False) -> bool:
+        """Determine if a node should be shown based on depth settings.
+
+        Depth is calculated by counting dots in the qualified name:
+        - depth 0: top-level packages (no dots)
+        - depth 1: first level sub-packages (one dot)
+        - depth 2: second level sub-packages (two dots)
+
+        For classes, depth is measured from their containing module, excluding
+        the class name itself from the depth calculation.
+        """
+        # If no depth limit is set ==> show all nodes
+        if self.max_depth is None:
+            return True
+
+        # For classes, we want to measure depth from their containing module
+        if is_class:
+            # Get the module part (everything before the last dot)
+            last_dot = qualified_name.rfind(".")
+            if last_dot == -1:
+                module_path = ""
+            else:
+                module_path = qualified_name[:last_dot]
+
+            # Count module depth
+            module_depth = module_path.count(".")
+            return bool(module_depth <= self.max_depth)
+
+        # For packages/modules, count full depth
+        node_depth = qualified_name.count(".")
+        return bool(node_depth <= self.max_depth)
+
     def write_packages(self, diagram: PackageDiagram) -> None:
         """Write a package diagram."""
         module_info: dict[str, dict[str, int]] = {}
@@ -61,6 +94,10 @@ def write_packages(self, diagram: PackageDiagram) -> None:
         for module in sorted(diagram.modules(), key=lambda x: x.title):
             module.fig_id = module.node.qname()
 
+            # Filter nodes based on depth
+            if not self.should_show_node(module.fig_id):
+                continue
+
             if self.config.no_standalone and not any(
                 module in (rel.from_object, rel.to_object)
                 for rel in diagram.get_relationships("depends")
@@ -83,6 +120,10 @@ def write_packages(self, diagram: PackageDiagram) -> None:
             from_id = rel.from_object.fig_id
             to_id = rel.to_object.fig_id
 
+            # Filter nodes based on depth ==> skip if either source or target nodes is beyond the max depth
+            if not self.should_show_node(from_id) or not self.should_show_node(to_id):
+                continue
+
             self.printer.emit_edge(
                 from_id,
                 to_id,
@@ -96,6 +137,10 @@ def write_packages(self, diagram: PackageDiagram) -> None:
             from_id = rel.from_object.fig_id
             to_id = rel.to_object.fig_id
 
+            # Filter nodes based on depth ==> skip if either source or target nodes is beyond the max depth
+            if not self.should_show_node(from_id) or not self.should_show_node(to_id):
+                continue
+
             self.printer.emit_edge(
                 from_id,
                 to_id,
@@ -115,6 +160,11 @@ def write_classes(self, diagram: ClassDiagram) -> None:
         # sorted to get predictable (hence testable) results
         for obj in sorted(diagram.objects, key=lambda x: x.title):
             obj.fig_id = obj.node.qname()
+
+            # Filter class based on depth setting
+            if not self.should_show_node(obj.fig_id, is_class=True):
+                continue
+
             if self.config.no_standalone and not any(
                 obj in (rel.from_object, rel.to_object)
                 for rel_type in ("specialization", "association", "aggregation")
@@ -129,6 +179,12 @@ def write_classes(self, diagram: ClassDiagram) -> None:
             )
         # inheritance links
         for rel in diagram.get_relationships("specialization"):
+            # Filter nodes based on depth setting
+            if not self.should_show_node(
+                rel.from_object.fig_id, is_class=True
+            ) or not self.should_show_node(rel.to_object.fig_id, is_class=True):
+                continue
+
             self.printer.emit_edge(
                 rel.from_object.fig_id,
                 rel.to_object.fig_id,
@@ -137,6 +193,12 @@ def write_classes(self, diagram: ClassDiagram) -> None:
         associations: dict[str, set[str]] = defaultdict(set)
         # generate associations
         for rel in diagram.get_relationships("association"):
+            # Filter nodes based on depth setting
+            if not self.should_show_node(
+                rel.from_object.fig_id, is_class=True
+            ) or not self.should_show_node(rel.to_object.fig_id, is_class=True):
+                continue
+
             associations[rel.from_object.fig_id].add(rel.to_object.fig_id)
             self.printer.emit_edge(
                 rel.from_object.fig_id,
@@ -146,6 +208,12 @@ def write_classes(self, diagram: ClassDiagram) -> None:
             )
         # generate aggregations
         for rel in diagram.get_relationships("aggregation"):
+            # Filter nodes based on depth setting
+            if not self.should_show_node(
+                rel.from_object.fig_id, is_class=True
+            ) or not self.should_show_node(rel.to_object.fig_id, is_class=True):
+                continue
+
             if rel.to_object.fig_id in associations[rel.from_object.fig_id]:
                 continue
             self.printer.emit_edge(
diff --git a/pylint/testutils/pyreverse.py b/pylint/testutils/pyreverse.py
@@ -23,6 +23,7 @@ class PyreverseConfig(
     The default values correspond to the defaults of the options' parser.
     """
 
+    # pylint: disable=too-many-locals
     def __init__(
         self,
         *,
@@ -40,6 +41,7 @@ def __init__(
         output_format: str = "dot",
         colorized: bool = False,
         max_color_depth: int = 2,
+        max_depth: int | None = None,
         color_palette: tuple[str, ...] = DEFAULT_COLOR_PALETTE,
         ignore_list: tuple[str, ...] = tuple(),
         project: str = "",
@@ -62,12 +64,15 @@ def __init__(
         self.only_classnames = only_classnames
         self.output_format = output_format
         self.colorized = colorized
+        self.max_depth = max_depth
         self.max_color_depth = max_color_depth
         self.color_palette = color_palette
         self.ignore_list = ignore_list
         self.project = project
         self.output_directory = output_directory
 
+    # pylint: enable=too-many-locals
+
 
 class TestFileOptions(TypedDict):
     source_roots: list[str]
diff --git a/pylintrc b/pylintrc
@@ -393,7 +393,7 @@ missing-member-max-choices=1
 spelling-dict=
 
 # List of comma separated words that should not be checked.
-spelling-ignore-words=
+spelling-ignore-words=subpkg,MyClass
 
 # List of comma separated words that should be considered directives if they
 # appear and the beginning of a comment and should not be checked.
diff --git a/tests/pyreverse/conftest.py b/tests/pyreverse/conftest.py
@@ -67,6 +67,13 @@ def html_config() -> PyreverseConfig:
     )
 
 
+@pytest.fixture()
+def depth_limited_config(default_max_depth: int) -> PyreverseConfig:
+    return PyreverseConfig(
+        max_depth=default_max_depth,
+    )
+
+
 @pytest.fixture(scope="session")
 def get_project() -> GetProjectCallable:
     def _get_project(module: str, name: str | None = "No Name") -> Project:
diff --git a/tests/pyreverse/data/classes_depth_limited_0.dot b/tests/pyreverse/data/classes_depth_limited_0.dot
@@ -0,0 +1,4 @@
+digraph "classes_depth_limited_0" {
+rankdir=BT
+charset="utf-8"
+}
diff --git a/tests/pyreverse/data/classes_depth_limited_1.dot b/tests/pyreverse/data/classes_depth_limited_1.dot
@@ -0,0 +1,17 @@
+digraph "classes_depth_limited_1" {
+rankdir=BT
+charset="utf-8"
+"data.clientmodule_test.Ancestor" [color="black", fontcolor="black", label=<{Ancestor|attr : str<br ALIGN="LEFT"/>cls_member<br ALIGN="LEFT"/>|get_value()<br ALIGN="LEFT"/>set_value(value)<br ALIGN="LEFT"/>}>, shape="record", style="solid"];
+"data.suppliermodule_test.CustomException" [color="black", fontcolor="red", label=<{CustomException|<br ALIGN="LEFT"/>|}>, shape="record", style="solid"];
+"data.suppliermodule_test.DoNothing" [color="black", fontcolor="black", label=<{DoNothing|<br ALIGN="LEFT"/>|}>, shape="record", style="solid"];
+"data.suppliermodule_test.DoNothing2" [color="black", fontcolor="black", label=<{DoNothing2|<br ALIGN="LEFT"/>|}>, shape="record", style="solid"];
+"data.suppliermodule_test.DoSomething" [color="black", fontcolor="black", label=<{DoSomething|my_int : Optional[int]<br ALIGN="LEFT"/>my_int_2 : Optional[int]<br ALIGN="LEFT"/>my_string : str<br ALIGN="LEFT"/>|do_it(new_int: int): int<br ALIGN="LEFT"/>}>, shape="record", style="solid"];
+"data.suppliermodule_test.Interface" [color="black", fontcolor="black", label=<{Interface|<br ALIGN="LEFT"/>|<I>get_value</I>()<br ALIGN="LEFT"/><I>set_value</I>(value)<br ALIGN="LEFT"/>}>, shape="record", style="solid"];
+"data.nullable_pattern.NullablePatterns" [color="black", fontcolor="black", label=<{NullablePatterns|<br ALIGN="LEFT"/>|<I>return_nullable_1</I>(): int \| None<br ALIGN="LEFT"/><I>return_nullable_2</I>(): Optional[int]<br ALIGN="LEFT"/>}>, shape="record", style="solid"];
+"data.property_pattern.PropertyPatterns" [color="black", fontcolor="black", label=<{PropertyPatterns|prop1<br ALIGN="LEFT"/>prop2<br ALIGN="LEFT"/>|}>, shape="record", style="solid"];
+"data.clientmodule_test.Specialization" [color="black", fontcolor="black", label=<{Specialization|TYPE : str<br ALIGN="LEFT"/>relation<br ALIGN="LEFT"/>relation2<br ALIGN="LEFT"/>top : str<br ALIGN="LEFT"/>|from_value(value: int)<br ALIGN="LEFT"/>increment_value(): None<br ALIGN="LEFT"/>transform_value(value: int): int<br ALIGN="LEFT"/>}>, shape="record", style="solid"];
+"data.clientmodule_test.Specialization" -> "data.clientmodule_test.Ancestor" [arrowhead="empty", arrowtail="none"];
+"data.suppliermodule_test.DoNothing" -> "data.clientmodule_test.Ancestor" [arrowhead="diamond", arrowtail="none", fontcolor="green", label="cls_member", style="solid"];
+"data.suppliermodule_test.DoNothing" -> "data.clientmodule_test.Specialization" [arrowhead="diamond", arrowtail="none", fontcolor="green", label="relation", style="solid"];
+"data.suppliermodule_test.DoNothing2" -> "data.clientmodule_test.Specialization" [arrowhead="odiamond", arrowtail="none", fontcolor="green", label="relation2", style="solid"];
+}
diff --git a/tests/pyreverse/data/packages_depth_limited_0.dot b/tests/pyreverse/data/packages_depth_limited_0.dot
@@ -0,0 +1,5 @@
+digraph "packages_depth_limited_0" {
+rankdir=BT
+charset="utf-8"
+"data" [color="black", label=<data>, shape="box", style="solid"];
+}
diff --git a/tests/pyreverse/data/packages_depth_limited_1.dot b/tests/pyreverse/data/packages_depth_limited_1.dot
@@ -0,0 +1,10 @@
+digraph "packages_depth_limited_1" {
+rankdir=BT
+charset="utf-8"
+"data" [color="black", label=<data>, shape="box", style="solid"];
+"data.clientmodule_test" [color="black", label=<data.clientmodule_test>, shape="box", style="solid"];
+"data.nullable_pattern" [color="black", label=<data.nullable_pattern>, shape="box", style="solid"];
+"data.property_pattern" [color="black", label=<data.property_pattern>, shape="box", style="solid"];
+"data.suppliermodule_test" [color="black", label=<data.suppliermodule_test>, shape="box", style="solid"];
+"data.clientmodule_test" -> "data.suppliermodule_test" [arrowhead="open", arrowtail="none"];
+}
diff --git a/tests/pyreverse/test_writer.py b/tests/pyreverse/test_writer.py
@@ -160,6 +160,18 @@ def setup_html(
     yield from _setup(project, html_config, writer)
 
 
+@pytest.fixture()
+def setup_depth_limited(
+    depth_limited_config: PyreverseConfig, get_project: GetProjectCallable
+) -> Iterator[None]:
+    writer = DiagramWriter(depth_limited_config)
+
+    project = get_project(
+        TEST_DATA_DIR, name=f"depth_limited_{depth_limited_config.max_depth}"
+    )
+    yield from _setup(project, depth_limited_config, writer)
+
+
 def _setup(
     project: Project, config: PyreverseConfig, writer: DiagramWriter
 ) -> Iterator[None]:
@@ -220,6 +232,14 @@ def test_colorized_puml_files(generated_file: str) -> None:
     _assert_files_are_equal(generated_file)
 
 
+@pytest.mark.parametrize("default_max_depth", [0, 1])
+@pytest.mark.usefixtures("setup_depth_limited")
+def test_depth_limited_write(default_max_depth: int) -> None:
+    """Test package diagram generation with a depth limit of 1."""
+    _assert_files_are_equal(f"packages_depth_limited_{default_max_depth}.dot")
+    _assert_files_are_equal(f"classes_depth_limited_{default_max_depth}.dot")
+
+
 def _assert_files_are_equal(generated_file: str) -> None:
     expected_file = os.path.join(os.path.dirname(__file__), "data", generated_file)
     generated = _file_lines(generated_file)
@@ -257,3 +277,84 @@ def test_package_name_with_slash(default_config: PyreverseConfig) -> None:
     writer.write([obj])
 
     assert os.path.exists("test_package_name_with_slash_.dot")
+
+
+def test_should_show_node_no_depth_limit(default_config: PyreverseConfig) -> None:
+    """Test that nodes are shown when no depth limit is set."""
+    writer = DiagramWriter(default_config)
+    writer.max_depth = None
+
+    assert writer.should_show_node("pkg")
+    assert writer.should_show_node("pkg.subpkg")
+    assert writer.should_show_node("pkg.subpkg.module")
+    assert writer.should_show_node("pkg.subpkg.module.submodule")
+
+
+@pytest.mark.parametrize("max_depth", range(5))
+def test_should_show_node_with_depth_limit(
+    default_config: PyreverseConfig, max_depth: int
+) -> None:
+    """Test that nodes are filtered correctly when depth limit is set.
+
+    Depth counting is zero-based, determined by number of dots in path:
+    - 'pkg'                  -> depth 0 (0 dots)
+    - 'pkg.subpkg'           -> depth 1 (1 dot)
+    - 'pkg.subpkg.module'    -> depth 2 (2 dots)
+    - 'pkg.subpkg.module.submodule' -> depth 3 (3 dots)
+    """
+    writer = DiagramWriter(default_config)
+    print("max_depth:", max_depth)
+    writer.max_depth = max_depth
+
+    # Test cases for different depths
+    test_cases = [
+        "pkg",
+        "pkg.subpkg",
+        "pkg.subpkg.module",
+        "pkg.subpkg.module.submodule",
+    ]
+
+    # Test if nodes are shown based on their depth and max_depth setting
+    for i, path in enumerate(test_cases):
+        should_show = i <= max_depth
+        print(
+            f"Path {path} (depth {i}) with max_depth={max_depth} "
+            f"{'should show' if should_show else 'should not show'}:"
+            f"{writer.should_show_node(path, is_class=True)}"
+        )
+        assert writer.should_show_node(path) == should_show
+
+
+@pytest.mark.parametrize("max_depth", range(5))
+def test_should_show_node_classes(
+    default_config: PyreverseConfig, max_depth: int
+) -> None:
+    """Test class visibility based on their containing module depth.
+
+     Classes are filtered based on their containing module's depth:
+    - 'MyClass' -> depth 0 (no module)
+    - 'pkg.MyClass' -> depth 0 (module has no dots)
+    - 'pkg.subpkg.MyClass' -> depth 1 (module has 1 dot)
+    - 'pkg.subpkg.mod.MyClass' -> depth 2 (module has 2 dots)
+    """
+    writer = DiagramWriter(default_config)
+    print("max_depth:", max_depth)
+    writer.max_depth = max_depth
+
+    # Test cases for different depths
+    test_cases = [
+        "MyClass",
+        "pkg.MyClass",
+        "pkg.subpkg.MyClass",
+        "pkg.subpkg.mod.MyClass",
+    ]
+
+    # Test if nodes are shown based on their depth and max_depth setting
+    for i, path in enumerate(test_cases):
+        should_show = i - 1 <= max_depth  # Subtract 1 to account for the class name
+        print(
+            f"Path {path} (depth {i}) with max_depth={max_depth} "
+            f"{'should show' if should_show else 'should not show'}:"
+            f"{writer.should_show_node(path, is_class=True)}"
+        )
+        assert writer.should_show_node(path, is_class=True) == should_show
