From 97a00ad6e1d740d1345ef309a085859bcdde1fef Mon Sep 17 00:00:00 2001 From: Vincent Koppen Date: Mon, 1 Jun 2026 09:56:48 +0200 Subject: [PATCH 1/3] fix: don't change graph on incorrect call to tmp_remove_nodes Signed-off-by: Vincent Koppen --- .../_core/model/graphs/models/base.py | 23 ++++--- tests/unit/model/graphs/test_graph_model.py | 62 ++++++++++++------- 2 files changed, 52 insertions(+), 33 deletions(-) diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py index 77da3c9a..f9926df4 100644 --- a/src/power_grid_model_ds/_core/model/graphs/models/base.py +++ b/src/power_grid_model_ds/_core/model/graphs/models/base.py @@ -212,16 +212,21 @@ def tmp_remove_nodes(self, nodes: list[int]) -> Generator: considering certain nodes. """ edge_list = [] - for node in nodes: - edge_list += list(self.in_branches(node)) - self.delete_node(node) + node_list = [] - yield - - for node in nodes: - self.add_node(int(node)) # convert to int to avoid type issues when input is e.g. a numpy array - for source, target in edge_list: - self.add_branch(source, target) + try: + for node in nodes: + edge_list += list(self.in_branches(node)) + + self.delete_node(node) + node_list.append(node) + + yield + finally: + for node in node_list: + self.add_node(int(node)) + for source, target in edge_list: + self.add_branch(source, target) def get_shortest_path(self, ext_start_node_id: int, ext_end_node_id: int) -> tuple[list[int], int]: """Calculate the shortest path between two nodes diff --git a/tests/unit/model/graphs/test_graph_model.py b/tests/unit/model/graphs/test_graph_model.py index e5a4fd4b..34998045 100644 --- a/tests/unit/model/graphs/test_graph_model.py +++ b/tests/unit/model/graphs/test_graph_model.py @@ -123,7 +123,7 @@ def test_internal_ids_after_node_deletion(self, graph: BaseGraphModel): assert graph._has_node(internal_id_1) assert graph._has_node(internal_id_2) - # now delete node 2, this can change the internal mapping + # now delete node 2, this can change the internal mapping% graph.delete_node(2) internal_id_0 = graph.external_to_internal(1) @@ -143,39 +143,53 @@ def test_graph_in_branches(self, graph: BaseGraphModel): assert list(graph.in_branches(2)) == [(1, 2), (1, 2), (1, 2)] -def test_tmp_remove_nodes(graph_with_2_routes: BaseGraphModel) -> None: - graph = graph_with_2_routes +class TestTmpRemoveNodes: + def test_tmp_remove_nodes(self, graph_with_2_routes: BaseGraphModel) -> None: + graph = graph_with_2_routes + + assert graph.nr_branches == 4 + + # add parallel branches to test whether they are restored correctly + graph.add_branch(1, 5) + graph.add_branch(5, 1) - assert graph.nr_branches == 4 + assert graph.nr_nodes == 5 + assert graph.nr_branches == 6 - # add parallel branches to test whether they are restored correctly - graph.add_branch(1, 5) - graph.add_branch(5, 1) + before_sets = [frozenset(branch) for branch in graph.all_branches] + counter_before = Counter(before_sets) - assert graph.nr_nodes == 5 - assert graph.nr_branches == 6 + with graph.tmp_remove_nodes([1, 2]): + assert graph.nr_nodes == 3 + assert list(graph.all_branches) == [(5, 4)] - before_sets = [frozenset(branch) for branch in graph.all_branches] - counter_before = Counter(before_sets) + assert graph.nr_nodes == 5 + assert graph.nr_branches == 6 - with graph.tmp_remove_nodes([1, 2]): - assert graph.nr_nodes == 3 - assert list(graph.all_branches) == [(5, 4)] + after_sets = [frozenset(branch) for branch in graph.all_branches] + counter_after = Counter(after_sets) + assert counter_before == counter_after - assert graph.nr_nodes == 5 - assert graph.nr_branches == 6 + def test_tmp_remove_nodes_array_input(self, graph_with_2_routes: BaseGraphModel) -> None: + with graph_with_2_routes.tmp_remove_nodes(np.array([1, 2])): # type: ignore[arg-type] + pass - after_sets = [frozenset(branch) for branch in graph.all_branches] - counter_after = Counter(after_sets) - assert counter_before == counter_after + # check that the external ids are still all integers instead of e.g. np.int + assert all([isinstance(e_id, int) for e_id in graph_with_2_routes.external_ids]) + def test_invalid_tmp_remove_nodes(self, graph_with_2_routes: BaseGraphModel) -> None: + original_graph = deepcopy(graph_with_2_routes) + assert graph_with_2_routes.nr_nodes == 5 + assert graph_with_2_routes.nr_branches == 4 -def test_tmp_remove_nodes_array_input(graph_with_2_routes: BaseGraphModel) -> None: - with graph_with_2_routes.tmp_remove_nodes(np.array([1, 2])): # type: ignore[arg-type] - pass + # When we remove node 1 and then an non-existing node that crashes the process + with pytest.raises(MissingNodeError), graph_with_2_routes.tmp_remove_nodes([1, 99]): + pass - # check that the external ids are still all integers instead of e.g. np.int - assert all([isinstance(e_id, int) for e_id in graph_with_2_routes.external_ids]) + # The remaining graph object should still contain the same nodes and edges. + assert graph_with_2_routes.nr_nodes == 5 + assert graph_with_2_routes.nr_branches == 4 + assert graph_with_2_routes == original_graph def test_get_components(graph_with_2_routes: BaseGraphModel): From ea791c986697bf50acb76498a49fc768bc38fb3c Mon Sep 17 00:00:00 2001 From: Vincent Koppen Date: Mon, 1 Jun 2026 09:58:52 +0200 Subject: [PATCH 2/3] cleanup Signed-off-by: Vincent Koppen --- tests/unit/model/graphs/test_graph_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/model/graphs/test_graph_model.py b/tests/unit/model/graphs/test_graph_model.py index 34998045..c9b09e3c 100644 --- a/tests/unit/model/graphs/test_graph_model.py +++ b/tests/unit/model/graphs/test_graph_model.py @@ -123,7 +123,7 @@ def test_internal_ids_after_node_deletion(self, graph: BaseGraphModel): assert graph._has_node(internal_id_1) assert graph._has_node(internal_id_2) - # now delete node 2, this can change the internal mapping% + # now delete node 2, this can change the internal mapping graph.delete_node(2) internal_id_0 = graph.external_to_internal(1) From 0306acc28b0461ef97d1d65a8b3a6b5bb4c246e3 Mon Sep 17 00:00:00 2001 From: Vincent Koppen Date: Mon, 1 Jun 2026 10:05:38 +0200 Subject: [PATCH 3/3] keep as is Signed-off-by: Vincent Koppen --- src/power_grid_model_ds/_core/model/graphs/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py index f9926df4..f87a335a 100644 --- a/src/power_grid_model_ds/_core/model/graphs/models/base.py +++ b/src/power_grid_model_ds/_core/model/graphs/models/base.py @@ -224,7 +224,7 @@ def tmp_remove_nodes(self, nodes: list[int]) -> Generator: yield finally: for node in node_list: - self.add_node(int(node)) + self.add_node(int(node)) # convert to int to avoid type issues when input is e.g. a numpy array for source, target in edge_list: self.add_branch(source, target)