Fix: prevent type change when using tmp_remove_nodes (#60)

* Fix: prevent type change in graph.external_ids when using tmp_remove_nodes_array

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

* provide correct type to tmp_remove_nodes

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

---------

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

Author: Thijss

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

@@ -208,7 +208,7 @@ def tmp_remove_nodes(self, nodes: list[int]) -> Generator:
         yield
 
         for node in nodes:
-            self.add_node(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)
 

src/power_grid_model_ds/_core/model/grids/helpers.py

@@ -46,7 +46,7 @@ def set_feeder_ids(grid: "Grid"):
     """
     _reset_feeder_ids(grid)
     feeder_node_ids = grid.node.filter(node_type=NodeType.SUBSTATION_NODE)["id"]
-    with grid.graphs.active_graph.tmp_remove_nodes(feeder_node_ids):
+    with grid.graphs.active_graph.tmp_remove_nodes(feeder_node_ids.tolist()):
         components = grid.graphs.active_graph.get_components()
     for component_node_ids in components:
         component_branches = _get_active_component_branches(grid, component_node_ids)

tests/unit/model/graphs/test_graph_model.py

@@ -171,6 +171,14 @@ def test_tmp_remove_nodes(graph_with_2_routes: BaseGraphModel) -> None:
     assert counter_before == counter_after
 
 
+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
+
+    # 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_get_components(graph_with_2_routes: BaseGraphModel):
     graph = graph_with_2_routes
     graph.add_node(99)