make .in_edges a public method

Signed-off-by: Thijs Baaijen <[email protected]>

src/power_grid_model_ds/_core/model/graphs/models/base.py

@@ -73,6 +73,14 @@ def has_node(self, node_id: int) -> bool:
 
         return self._has_node(node_id=internal_node_id)
 
+    def in_edges(self, node_id: int) -> Generator[tuple[int, int], None, None]:
+        """Return all edges a node occurs in."""
+        int_node_id = self.external_to_internal(node_id)
+        internal_edges = self._in_edges(int_node_id=int_node_id)
+        return (
+            (self.internal_to_external(source), self.internal_to_external(target)) for source, target in internal_edges
+        )
+
     def add_node(self, ext_node_id: int, raise_on_fail: bool = True) -> None:
         """Add a node to the graph."""
         if self.has_node(ext_node_id):
@@ -187,11 +195,7 @@ def tmp_remove_nodes(self, nodes: list[int]) -> Generator:
         edge_list = []
         for node in nodes:
             internal_node = self.external_to_internal(node)
-            node_edges = [
-                (self.internal_to_external(source), self.internal_to_external(target))
-                for source, target in self._in_edges(internal_node)
-            ]
-            edge_list += node_edges
+            edge_list += list(self.in_edges(node))
             self._delete_node(internal_node)
         yield
 
@@ -307,10 +311,10 @@ def _branch_is_relevant(self, branch: BranchArray) -> bool:
         return True
 
     @abstractmethod
-    def _in_edges(self, internal_node: int) -> list[tuple[int, int]]:
+    def _in_edges(self, int_node_id: int) -> Generator[tuple[int, int], None, None]:
         """Return all edges a node occurs in.
-
-        Return a list of tuples with the source and target node id. These are internal node ids.
+        Return a list of tuples with the source and target node id.
+        These are internal node ids.
         """
 
     @abstractmethod

src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py

@@ -100,8 +100,8 @@ def _get_connected(self, node_id: int, nodes_to_ignore: list[int], inclusive: bo
 
         return connected_nodes
 
-    def _in_edges(self, internal_node: int) -> list[tuple[int, int]]:
-        return [(source, target) for source, target, _ in self._graph.in_edges(internal_node)]
+    def _in_edges(self, int_node_id: int) -> Generator[tuple[int, int], None, None]:
+        return ((source, target) for source, target, _ in self._graph.in_edges(int_node_id))
 
     def _find_fundamental_cycles(self) -> list[list[int]]:
         """Find all fundamental cycles in the graph using Rustworkx.

tests/unit/model/graphs/test_graph_model.py

@@ -57,6 +57,17 @@ def test_graph_all_branches_parallel(graph):
     assert [(1, 2), (1, 2), (2, 1)] == list(graph.all_branches)
 
 
+def test_graph_in_edges(graph):
+    graph.add_node(1)
+    graph.add_node(2)
+    graph.add_branch(1, 2)
+    graph.add_branch(1, 2)
+    graph.add_branch(2, 1)
+
+    assert [(2, 1), (2, 1), (2, 1)] == list(graph.in_edges(1))
+    assert [(1, 2), (1, 2), (1, 2)] == list(graph.in_edges(2))
+
+
 def test_graph_delete_branch(graph):
     """Test whether a branch is deleted correctly"""
     graph.add_node(1)