Apply suggestions from code review

Co-authored-by: Vincent Koppen <[email protected]>
Signed-off-by: Thijs Baaijen <[email protected]>

Author: Thijss

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

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MPL-2.0
 
 from abc import ABC, abstractmethod
+from typing import Generator
 
 import numpy as np
 from numpy._typing import NDArray
@@ -35,11 +36,14 @@ def nr_branches(self):
         """Returns the number of branches in the graph"""
 
     @property
-    @abstractmethod
-    def all_branches(self) -> list[frozenset[int]]:
+    def all_branches(self) -> Generator[tuple[int, int], None, None]:
         """Returns all branches in the graph as a list of node pairs (frozensets).
         Warning: Depending on graph engine, performance could be slow for large graphs
         """
+        return (
+            (self.internal_to_external(source), self.internal_to_external(target))
+            for source, target in self._all_branches()
+        )
 
     @abstractmethod
     def external_to_internal(self, ext_node_id: int) -> int:
@@ -314,6 +318,9 @@ def _get_components(self, substation_nodes: list[int]) -> list[list[int]]: ...
     @abstractmethod
     def _find_fundamental_cycles(self) -> list[list[int]]: ...
 
+    @abstractmethod
+    def _all_branches(self) -> Generator[tuple[int, int], None, None]: ...
+
 
 def _get_branch3_branches(branch3: Branch3Array) -> BranchArray:
     node_1 = branch3.node_1.item()

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

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MPL-2.0
 
 import logging
+from typing import Generator
 
 import rustworkx as rx
 from rustworkx import NoEdgeBetweenNodes
@@ -32,15 +33,6 @@ def nr_nodes(self):
     def nr_branches(self):
         return self._graph.num_edges()
 
-    @property
-    def all_branches(self) -> list[frozenset[int]]:
-        internal_branches = ((source, target) for source, target in self._graph.edge_list())
-        external_branches = [
-            frozenset([self.internal_to_external(source), self.internal_to_external(target)])
-            for source, target in internal_branches
-        ]
-        return external_branches
-
     @property
     def external_ids(self) -> list[int]:
         return list(self._external_to_internal.keys())
@@ -116,6 +108,9 @@ def _find_fundamental_cycles(self) -> list[list[int]]:
         """
         return find_fundamental_cycles_rustworkx(self._graph)
 
+    def _all_branches(self) -> Generator[tuple[int, int], None, None]:
+        return ((source, target) for source, target in self._graph.edge_list())
+
 
 class _NodeVisitor(BFSVisitor):
     def __init__(self, nodes_to_ignore: list[int]):

tests/unit/model/graphs/test_graph_model.py

@@ -42,7 +42,7 @@ def test_graph_all_branches(graph):
     graph.add_node(2)
     graph.add_branch(1, 2)
 
-    assert [{1, 2}] == graph.all_branches
+    assert [(1, 2)] == list(graph.all_branches)
 
 
 def test_graph_all_branches_parallel(graph):
@@ -52,7 +52,7 @@ def test_graph_all_branches_parallel(graph):
     graph.add_branch(1, 2)
     graph.add_branch(2, 1)
 
-    assert [{1, 2}, {1, 2}, {1, 2}] == graph.all_branches
+    assert [(1, 2), (1, 2), (2, 1)] == list(graph.all_branches)
 
 
 def test_graph_delete_branch(graph):