Refactor depth filtering logic

pylint/pyreverse/diadefslib.py

@@ -67,6 +67,16 @@ def _get_levels(self) -> tuple[int, int]:
         """Help function for search levels."""
         return self.anc_level, self.association_level
 
+    def _should_include_by_depth(self, node: nodes.NodeNG) -> bool:
+        """Check if a node should be included based on depth."""
+        # If max_depth is not set, include all nodes
+        if self.config.max_depth is None:
+            return True
+
+        # For other nodes, calculate depth based on their root module
+        module_depth = node.root().name.count(".")
+        return module_depth <= self.config.max_depth  # type: ignore[no-any-return]
+
     def show_node(self, node: nodes.ClassDef) -> bool:
         """Determine if node should be shown based on config."""
         if node.root().name == "builtins":
@@ -75,7 +85,8 @@ def show_node(self, node: nodes.ClassDef) -> bool:
         if is_stdlib_module(node.root().name):
             return self.config.show_stdlib  # type: ignore[no-any-return]
 
-        return True
+        # Filter node by depth
+        return self._should_include_by_depth(node)
 
     def add_class(self, node: nodes.ClassDef) -> None:
         """Visit one class and add it to diagram."""
@@ -163,7 +174,7 @@ def visit_module(self, node: nodes.Module) -> None:
 
         add this class to the package diagram definition
         """
-        if self.pkgdiagram:
+        if self.pkgdiagram and self._should_include_by_depth(node):
             self.linker.visit(node)
             self.pkgdiagram.add_object(node.name, node)
 
@@ -177,7 +188,7 @@ def visit_classdef(self, node: nodes.ClassDef) -> None:
 
     def visit_importfrom(self, node: nodes.ImportFrom) -> None:
         """Visit astroid.ImportFrom  and catch modules for package diagram."""
-        if self.pkgdiagram:
+        if self.pkgdiagram and self._should_include_by_depth(node):
             self.pkgdiagram.add_from_depend(node, node.modname)
 
 

pylint/pyreverse/writer.py

@@ -54,38 +54,6 @@ 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]] = {}
@@ -94,10 +62,6 @@ 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")
@@ -120,10 +84,6 @@ 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,
@@ -137,10 +97,6 @@ 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,
@@ -161,10 +117,6 @@ def write_classes(self, diagram: ClassDiagram) -> None:
         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")
@@ -179,12 +131,6 @@ 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,
@@ -193,12 +139,6 @@ 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,
@@ -208,12 +148,6 @@ 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(

tests/pyreverse/test_diadefs.py

@@ -8,8 +8,9 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from pathlib import Path
+from unittest.mock import Mock
 
 import pytest
 from astroid import extract_node, nodes
@@ -58,6 +59,24 @@ def PROJECT(get_project: GetProjectCallable) -> Iterator[Project]:
         yield get_project("data")
 
 
+# Fixture for creating mocked nodes
+@pytest.fixture
+def mock_node() -> Callable[[str], Mock]:
+    def _mock_node(module_path: str) -> Mock:
+        """Create a mocked node with a given module path."""
+        node = Mock()
+        node.root.return_value.name = module_path
+        return node
+
+    return _mock_node
+
+
+# Fixture for the DiaDefGenerator
+@pytest.fixture
+def generator(HANDLER: DiadefsHandler, PROJECT: Project) -> DiaDefGenerator:
+    return DiaDefGenerator(Linker(PROJECT), HANDLER)
+
+
 def test_option_values(
     default_config: PyreverseConfig, HANDLER: DiadefsHandler, PROJECT: Project
 ) -> None:
@@ -257,3 +276,56 @@ def test_regression_dataclasses_inference(
     special = "regrtest_data.dataclasses_pyreverse.InventoryItem"
     cd = cdg.class_diagram(path, special)
     assert cd.title == special
+
+
+# Test for no depth limit
+def test_should_include_by_depth_no_limit(
+    generator: DiaDefGenerator, mock_node: Mock
+) -> None:
+    """Test that nodes are included when no depth limit is set."""
+    generator.config.max_depth = None
+
+    # Create mocked nodes with different depths
+    node1 = mock_node("pkg")  # Depth 0
+    node2 = mock_node("pkg.subpkg")  # Depth 1
+    node3 = mock_node("pkg.subpkg.module")  # Depth 2
+
+    # All nodes should be included when max_depth is None
+    assert generator._should_include_by_depth(node1)
+    assert generator._should_include_by_depth(node2)
+    assert generator._should_include_by_depth(node3)
+
+
+@pytest.mark.parametrize("max_depth", range(5))
+def test_should_include_by_depth_with_limit(
+    generator: DiaDefGenerator, mock_node: Mock, 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)
+    """
+    # Set generator config
+    generator.config.max_depth = max_depth
+
+    # Test cases for different depths
+    test_cases = [
+        "pkg",
+        "pkg.subpkg",
+        "pkg.subpkg.module",
+        "pkg.subpkg.module.submodule",
+    ]
+    nodes = [mock_node(path) for path in test_cases]
+
+    # Test if nodes are shown based on their depth and max_depth setting
+    for i, node in enumerate(nodes):
+        should_show = i <= max_depth
+        print(
+            f"Node {node.root.return_value.name} (depth {i}) with max_depth={max_depth} "
+            f"{'should show' if should_show else 'should not show'}:"
+            f"{generator._should_include_by_depth(node)}"
+        )
+        assert generator._should_include_by_depth(node) == should_show

tests/pyreverse/test_writer.py

@@ -277,84 +277,3 @@ 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