Update writer to only show modules below max depth setting

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,15 @@ def write(self, diadefs: Iterable[ClassDiagram | PackageDiagram]) -> None:
                 self.write_classes(diagram)
             self.save()
 
+    def should_show_node(self, qualified_name: str) -> bool:
+        """Determine if a node should be shown based on depth settings."""
+        # If no depth limit is set ==> show all nodes
+        if self.max_depth is None:
+            return True
+        # Count dots to determine depth, subtract 1 since root counts as depth 0
+        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 +71,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 +97,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 +114,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 +137,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):
+                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 +156,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
+            ) or not self.should_show_node(rel.to_object.fig_id):
+                continue
+
             self.printer.emit_edge(
                 rel.from_object.fig_id,
                 rel.to_object.fig_id,
@@ -137,6 +170,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
+            ) or not self.should_show_node(rel.to_object.fig_id):
+                continue
+
             associations[rel.from_object.fig_id].add(rel.to_object.fig_id)
             self.printer.emit_edge(
                 rel.from_object.fig_id,
@@ -146,6 +185,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
+            ) or not self.should_show_node(rel.to_object.fig_id):
+                continue
+
             if rel.to_object.fig_id in associations[rel.from_object.fig_id]:
                 continue
             self.printer.emit_edge(