Skip to content

Commit 369498b

Browse files
committed
make .in_edges a public method
Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
1 parent 4cb80a4 commit 369498b

File tree

3 files changed

+25
-10
lines changed

3 files changed

+25
-10
lines changed

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

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,14 @@ def has_node(self, node_id: int) -> bool:
7373

7474
return self._has_node(node_id=internal_node_id)
7575

76+
def in_edges(self, node_id: int) -> Generator[tuple[int, int], None, None]:
77+
"""Return all edges a node occurs in."""
78+
int_node_id = self.external_to_internal(node_id)
79+
internal_edges = self._in_edges(int_node_id=int_node_id)
80+
return (
81+
(self.internal_to_external(source), self.internal_to_external(target)) for source, target in internal_edges
82+
)
83+
7684
def add_node(self, ext_node_id: int, raise_on_fail: bool = True) -> None:
7785
"""Add a node to the graph."""
7886
if self.has_node(ext_node_id):
@@ -187,11 +195,7 @@ def tmp_remove_nodes(self, nodes: list[int]) -> Generator:
187195
edge_list = []
188196
for node in nodes:
189197
internal_node = self.external_to_internal(node)
190-
node_edges = [
191-
(self.internal_to_external(source), self.internal_to_external(target))
192-
for source, target in self._in_edges(internal_node)
193-
]
194-
edge_list += node_edges
198+
edge_list += list(self.in_edges(node))
195199
self._delete_node(internal_node)
196200
yield
197201

@@ -307,10 +311,10 @@ def _branch_is_relevant(self, branch: BranchArray) -> bool:
307311
return True
308312

309313
@abstractmethod
310-
def _in_edges(self, internal_node: int) -> list[tuple[int, int]]:
314+
def _in_edges(self, int_node_id: int) -> Generator[tuple[int, int], None, None]:
311315
"""Return all edges a node occurs in.
312-
313-
Return a list of tuples with the source and target node id. These are internal node ids.
316+
Return a list of tuples with the source and target node id.
317+
These are internal node ids.
314318
"""
315319

316320
@abstractmethod

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,8 @@ def _get_connected(self, node_id: int, nodes_to_ignore: list[int], inclusive: bo
100100

101101
return connected_nodes
102102

103-
def _in_edges(self, internal_node: int) -> list[tuple[int, int]]:
104-
return [(source, target) for source, target, _ in self._graph.in_edges(internal_node)]
103+
def _in_edges(self, int_node_id: int) -> Generator[tuple[int, int], None, None]:
104+
return ((source, target) for source, target, _ in self._graph.in_edges(int_node_id))
105105

106106
def _find_fundamental_cycles(self) -> list[list[int]]:
107107
"""Find all fundamental cycles in the graph using Rustworkx.

tests/unit/model/graphs/test_graph_model.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,17 @@ def test_graph_all_branches_parallel(graph):
5757
assert [(1, 2), (1, 2), (2, 1)] == list(graph.all_branches)
5858

5959

60+
def test_graph_in_edges(graph):
61+
graph.add_node(1)
62+
graph.add_node(2)
63+
graph.add_branch(1, 2)
64+
graph.add_branch(1, 2)
65+
graph.add_branch(2, 1)
66+
67+
assert [(2, 1), (2, 1), (2, 1)] == list(graph.in_edges(1))
68+
assert [(1, 2), (1, 2), (1, 2)] == list(graph.in_edges(2))
69+
70+
6071
def test_graph_delete_branch(graph):
6172
"""Test whether a branch is deleted correctly"""
6273
graph.add_node(1)

0 commit comments

Comments
 (0)