Depth based filters

src/codoc/service/filters/__init__.py

@@ -6,6 +6,7 @@
 """
 from .class_diagram import class_diagram_filter
 from .children_based import get_children_of
+from .depth_based import get_depth_based_filter
 from .type_exclusion_filter import exclude_classes, exclude_functions, exclude_modules
 
 
@@ -15,4 +16,5 @@
     "exclude_functions",
     "exclude_modules",
     "get_children_of",
+    "get_depth_based_filter",
 ]

src/codoc/service/filters/children_based.py

@@ -3,6 +3,10 @@
 from codoc.domain.model import Graph, Node, NodeId, Dependency
 from codoc.domain.helpers import get_node, get_children, set_parent
 from .types import FilterType
+from .helpers import (
+    is_both_edges_of_edge_in_list_of_nodes,
+    is_either_edges_of_edge_in_list_of_nodes,
+)
 
 
 def get_children_of(identifier: str, keep_external_nodes: bool = False) -> FilterType:
@@ -96,12 +100,10 @@ def is_node_accepted(node: Node, graph: Graph, allowed_identifier: NodeId) -> bo
 def is_edge_accepted(
     edge: Dependency, node_identifiers: Set[str], keep_external_nodes: bool
 ) -> bool:
-    is_from_node_internal = edge.from_node in node_identifiers
-    is_to_node_internal = edge.to_node in node_identifiers
     if keep_external_nodes:
-        return is_from_node_internal or is_to_node_internal
+        return is_either_edges_of_edge_in_list_of_nodes(edge, node_identifiers)
     else:
-        return is_from_node_internal and is_to_node_internal
+        return is_both_edges_of_edge_in_list_of_nodes(edge, node_identifiers)
 
 
 def is_node_in_edges(node: Node, edges: Set[Dependency], graph: Graph) -> bool:

src/codoc/service/filters/depth_based.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+"""
+These filters look at the depth of a graph, and
+return new graphs that are based on the so-called depth of a given node.
+
+With "depth" we mean how close to the root of a graph it is, i.e
+a depth=1 would only take the top-level modules, where depth 2
+would take all nodes that are direct children of top-level modules.
+
+Depth can be found by recursively checking if the parent has a parent,
+and if yes check that parent too, and return how many steps one has to
+ progress before returning a null value.
+"""
+
+from codoc.domain.helpers import get_node
+from codoc.domain.model import Graph, Node
+
+from .types import FilterType
+
+from .helpers import get_edges_where_both_ends_are_in_nodes
+
+
+def get_depth_based_filter(depth: int) -> FilterType:
+    def depth_based_filter(graph: Graph) -> Graph:
+        nodes = set(
+            node for node in graph.nodes if get_depth_of_node(node, graph) <= depth
+        )
+        edges = get_edges_where_both_ends_are_in_nodes(graph.edges, nodes)
+        return Graph(nodes=nodes, edges=edges)
+
+    return depth_based_filter
+
+
+def get_depth_of_node(node: Node, graph: Graph) -> int:
+    parent_id = node.parent_identifier
+    if parent_id is None:
+        return 1
+    parent = parent_id = get_node(parent_id, graph)
+    return get_depth_of_node(parent, graph) + 1

src/codoc/service/filters/helpers.py

@@ -1,8 +1,35 @@
 #!/usr/bin/env python3
-from codoc.domain.model import Node
+from typing import Set
+
+from codoc.domain.model import Dependency, Node
 
 
 def node_without_parent(node: Node) -> Node:
     kwargs = node.__dict__
     kwargs["parent_identifier"] = None
     return Node(**kwargs)
+
+
+def get_edges_where_both_ends_are_in_nodes(edges: Set[Dependency], nodes: Set[Node]):
+    node_identifiers = set(node.identifier for node in nodes)
+    return set(
+        edge
+        for edge in edges
+        if is_both_edges_of_edge_in_list_of_nodes(edge, node_identifiers)
+    )
+
+
+def is_both_edges_of_edge_in_list_of_nodes(
+    edge: Dependency, node_identifiers: Set[str]
+) -> bool:
+    is_from_node_internal = edge.from_node in node_identifiers
+    is_to_node_internal = edge.to_node in node_identifiers
+    return is_from_node_internal and is_to_node_internal
+
+
+def is_either_edges_of_edge_in_list_of_nodes(
+    edge: Dependency, node_identifiers: Set[str]
+) -> bool:
+    is_from_node_internal = edge.from_node in node_identifiers
+    is_to_node_internal = edge.to_node in node_identifiers
+    return is_from_node_internal or is_to_node_internal

tests/unit/test_filters.py

@@ -2,7 +2,8 @@
 import pytest
 
 from codoc.service import filters
-from codoc.domain.model import NodeType
+from codoc.service.filters import depth_based
+from codoc.domain.model import NodeType, Graph
 from codoc.domain.helpers import contains_node, contains_dependency_between
 
 
@@ -20,6 +21,60 @@ def test_matches_snapshot(filter_function, test_graph, assert_match_snap):
     assert_match_snap(filtered_graph)
 
 
+class TestDepthBasedFilter:
+    def test_returns_filter(self, graph):
+        assert type(filters.get_depth_based_filter(1)(graph)) is Graph
+
+    def test_depth_zero_is_empty(self, graph, parent):
+        filtered_graph = filters.get_depth_based_filter(0)(graph)
+        assert len(filtered_graph.nodes) == 0 and len(filtered_graph.edges) == 0
+
+    def test_depth_one_only_has_root(self, graph, parent):
+        assert filters.get_depth_based_filter(1)(graph).nodes == {parent}
+
+    class TestGetDepthOfNode:
+        def test_depth_of_parent(self, parent, get_depth_of):
+            assert get_depth_of(parent) == 1
+
+        def test_depth_of_child(self, child, get_depth_of):
+            assert get_depth_of(child) == 2
+
+        def test_depth_of_child_2(self, child_2, get_depth_of):
+            assert get_depth_of(child_2) == 2
+
+        @pytest.fixture()
+        def get_depth_of(self, graph):
+            return lambda node: depth_based.get_depth_of_node(node, graph)
+
+    def test_depth_zero_has_no_edges(self, graph, parent):
+        assert len(filters.get_depth_based_filter(0)(graph).edges) == 0
+
+    def test_keep_only_relevant_edges(self, create_graph, create_edge, create_node):
+        nodeA = create_node(identifier="A")
+        nodeB = create_node(identifier="B")
+        nodeC = create_node(identifier="C", parent=nodeB)
+        edge1 = create_edge(nodeA, nodeB)
+        edge2 = create_edge(nodeA, nodeC)
+        graph = create_graph(nodes={nodeA, nodeB, nodeC}, edges={edge1, edge2})
+        assert filters.get_depth_based_filter(1)(graph).edges == {edge1}
+
+    @pytest.fixture()
+    def graph(self, create_graph, child, child_2, parent):
+        return create_graph(nodes={child_2, child, parent})
+
+    @pytest.fixture()
+    def child_2(self, create_node, parent):
+        return create_node(parent=parent, identifier="child")
+
+    @pytest.fixture()
+    def child(self, create_node, parent):
+        return create_node(parent=parent, identifier="child")
+
+    @pytest.fixture()
+    def parent(self, create_node):
+        return create_node(identifier="parent")
+
+
 class TestGetChildrenOfFilter:
     def test_matches_snapshot(self, assert_match_snap, filtered_graph):
         assert_match_snap(filtered_graph)
@@ -46,7 +101,7 @@ def filtered_graph(self, parent, child, graph):
 
         @pytest.fixture()
         def graph(self, create_graph, child, parent):
-            return create_graph(nodes=[child, parent])
+            return create_graph(nodes={child, parent})
 
         @pytest.fixture()
         def child(self, create_node, parent):