chore: add typing to tests

Signed-off-by: jaapschoutenalliander <[email protected]>

Author: jaapschoutenalliander

src/power_grid_model_ds/_core/model/containers/grid_protocol.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MPL-2.0
 """This file contains the Grid protocol defining the minimal arrays contained in a grid"""
 
+from abc import abstractmethod
 from typing import Protocol
 
 from power_grid_model_ds._core.model.arrays import (
@@ -18,5 +19,13 @@ class MinimalGridArrays(Protocol):
 
     node: NodeArray
     three_winding_transformer: ThreeWindingTransformerArray
-    branches: BranchArray
-    branch_arrays: list[BranchArray]
+
+    @property
+    @abstractmethod
+    def branches(self) -> BranchArray:
+        """Converts all branch arrays into a single BranchArray."""
+
+    @property
+    @abstractmethod
+    def branch_arrays(self) -> list[BranchArray]:
+        """Returns all branch arrays"""

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

@@ -35,6 +35,7 @@
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 from power_grid_model_ds._core.model.arrays.base.errors import RecordDoesNotExist
 from power_grid_model_ds._core.model.containers.base import FancyArrayContainer
+from power_grid_model_ds._core.model.containers.grid_protocol import MinimalGridArrays
 from power_grid_model_ds._core.model.enums.nodes import NodeType
 from power_grid_model_ds._core.model.graphs.container import GraphContainer
 from power_grid_model_ds._core.model.graphs.models import RustworkxGraphModel
@@ -51,7 +52,7 @@
 
 
 @dataclass
-class Grid(FancyArrayContainer):
+class Grid(FancyArrayContainer, MinimalGridArrays):
     """Grid object containing the entire network and interface to interact with it.
 
     Examples:

tests/conftest.py

@@ -64,12 +64,12 @@ def fancy_test_array():
 
 
 @pytest.fixture
-def basic_grid(grid):
+def basic_grid(grid: Grid):
     yield build_basic_grid(grid)
 
 
 @pytest.fixture
-def grid_with_3wt(grid):
+def grid_with_3wt(grid: Grid):
     yield build_basic_grid_with_three_winding(grid)
 
 

tests/integration/loadflow/test_power_grid_model.py

@@ -60,7 +60,7 @@ def test_load_flow_on_random():
     assert all(output["line"]["i_from"] > 0)
 
 
-def test_load_flow(grid):
+def test_load_flow(grid: Grid):
     """Tests the load flow on a test grid with 2 nodes"""
     nodes = NodeArray.zeros(2)
     nodes.id = [0, 1]
@@ -105,7 +105,7 @@ def test_load_flow(grid):
     assert all(output["line"]["i_from"] > 0)
 
 
-def test_load_flow_with_transformer(grid):
+def test_load_flow_with_transformer(grid: Grid):
     """Tests the load flow on a test grid with 3 nodes and a trafo"""
     nodes = NodeArray.zeros(3)
     nodes.id = [0, 1, 2]
@@ -178,7 +178,7 @@ def test_load_flow_with_transformer(grid):
 
 # pylint: disable=too-many-statements
 # pylint: disable=duplicate-code
-def test_load_flow_with_three_winding_transformer(grid):
+def test_load_flow_with_three_winding_transformer(grid: Grid):
     """Tests the load flow on a test grid with 3 nodes and a three winding trafo"""
     nodes = NodeArray.zeros(3)
     nodes.id = [0, 1, 2]
@@ -249,7 +249,7 @@ def test_load_flow_with_three_winding_transformer(grid):
     assert all(output["three_winding_transformer"]["loading"] > 0)
 
 
-def test_load_flow_with_link(grid):
+def test_load_flow_with_link(grid: Grid):
     """Tests the load flow on a test grid with 2 nodes and a link"""
     nodes = NodeArray.zeros(2)
     nodes.id = [0, 1]
@@ -293,7 +293,7 @@ def test_load_flow_with_link(grid):
     assert all(output["link"]["i_from"] > 0)
 
 
-def test_automatic_tap_regulator(grid):
+def test_automatic_tap_regulator(grid: Grid):
     """Test automatic tap regulator
 
     Network:

tests/unit/data_source/generator/test_grid_generators.py

@@ -14,6 +14,7 @@
 from power_grid_model_ds._core.data_source.generator.grid_generators import RadialGridGenerator
 from power_grid_model_ds._core.load_flow import PowerGridModelInterface
 from power_grid_model_ds._core.model.arrays import LineArray, NodeArray, SourceArray, SymLoadArray
+from power_grid_model_ds._core.model.graphs.models.base import BaseGraphModel
 from power_grid_model_ds._core.model.grids.base import Grid
 
 
@@ -37,14 +38,14 @@ def test_generate_random_grid():
     assert len(grid.line) == grid.graphs.complete_graph.nr_branches
 
 
-def test_graph_generate_random_grid_with_different_graph_engines(graph):
+def test_graph_generate_random_grid_with_different_graph_engines(graph: BaseGraphModel):
     """Generate a random grid with correct structure"""
     grid_generator = RadialGridGenerator(grid_class=Grid, graph_model=graph.__class__)
     grid = grid_generator.run(seed=0)
     assert isinstance(grid.graphs.active_graph, graph.__class__)
 
 
-def test_generate_random_nodes(grid):
+def test_generate_random_nodes(grid: Grid):
     """Generate random nodes"""
     node_generator = NodeGenerator(grid, seed=0)
     nodes, loads_low, loads_high = node_generator.run(amount=2)
@@ -64,7 +65,7 @@ def test_generate_random_nodes(grid):
     assert all(np.isin(loads_low.node, nodes.id))
 
 
-def test_generate_random_sources(grid):
+def test_generate_random_sources(grid: Grid):
     """Generate random sources"""
     source_generator = SourceGenerator(grid=grid, seed=0)
     nodes, sources = source_generator.run(amount=1)
@@ -81,7 +82,7 @@ def test_generate_random_sources(grid):
     assert all(np.isin(sources.node, nodes.id))
 
 
-def test_generate_random_lines(grid):
+def test_generate_random_lines(grid: Grid):
     """Generate random lines"""
     nodes = NodeArray.zeros(4)
     nodes.id = [0, 1, 2, 3]
@@ -112,7 +113,7 @@ def test_generate_random_lines(grid):
     assert all(np.isin(lines.to_node, nodes.id))
 
 
-def test_create_routes(grid):
+def test_create_routes(grid: Grid):
     """Generate new routes"""
     nodes = NodeArray.zeros(4)
     nodes.id = [0, 1, 2, 3]
@@ -146,7 +147,7 @@ def test_create_routes(grid):
     assert all(np.isin(line_generator.line_array.to_node, nodes.id))
 
 
-def test_determine_number_of_routes(grid):
+def test_determine_number_of_routes(grid: Grid):
     """Number of routes"""
     line_generator = LineGenerator(grid=grid, seed=0)
 
@@ -168,7 +169,7 @@ def test_determine_number_of_routes(grid):
     assert 3 == number_of_routes
 
 
-def test_connect_nodes(grid):
+def test_connect_nodes(grid: Grid):
     """Connect nodes"""
     nodes = NodeArray.zeros(4)
     nodes.id = [0, 1, 2, 3]
@@ -205,7 +206,7 @@ def test_connect_nodes(grid):
     assert 2 == len(line_generator.line_array)
 
 
-def test_create_nops(grid):
+def test_create_nops(grid: Grid):
     """Create normally open points"""
     nodes = NodeArray.zeros(4)
     nodes.id = [0, 1, 2, 3]

tests/unit/model/arrays/test_array.py

@@ -33,63 +33,63 @@ class _InheritedStrLengthArray(_CustomStrLengthArray):
     _str_lengths = {"extra_string": 100}
 
 
-def test_get_non_existing_attribute(fancy_test_array):
+def test_get_non_existing_attribute(fancy_test_array: FancyTestArray):
     with pytest.raises(AttributeError):
         # pylint: disable=pointless-statement
         fancy_test_array.non_existing_attribute  # noqa
 
 
-def test_for_loop(fancy_test_array):
+def test_for_loop(fancy_test_array: FancyTestArray):
     for row in fancy_test_array:
         assert row in fancy_test_array
         assert isinstance(row, FancyTestArray)
 
 
-def test_setattr(fancy_test_array):
+def test_setattr(fancy_test_array: FancyTestArray):
     assert_array_equal(fancy_test_array.id, [1, 2, 3])
     fancy_test_array.id = [9, 9, 9]
 
     assert_array_equal(fancy_test_array.id, [9, 9, 9])
     assert_array_equal(fancy_test_array.data["id"], [9, 9, 9])
 
 
-def test_prevent_delete_numpy_attribute(fancy_test_array):
+def test_prevent_delete_numpy_attribute(fancy_test_array: FancyTestArray):
     with pytest.raises(AttributeError):
         del fancy_test_array.size  # 'size' is a numpy attribute
 
 
-def test_getitem_array_one_column(fancy_test_array):
+def test_getitem_array_one_column(fancy_test_array: FancyTestArray):
     assert_array_equal(fancy_test_array["id"], [1, 2, 3])
 
 
-def test_getitem_array_multiple_columns(fancy_test_array):
+def test_getitem_array_multiple_columns(fancy_test_array: FancyTestArray):
     columns = ["id", "test_int", "test_float"]
     assert fancy_test_array.data[columns].tolist() == fancy_test_array[columns].tolist()
     assert_array_equal(fancy_test_array[columns].dtype.names, ("id", "test_int", "test_float"))
 
 
-def test_getitem_unique_multiple_columns(fancy_test_array):
+def test_getitem_unique_multiple_columns(fancy_test_array: FancyTestArray):
     columns = ["id", "test_int", "test_float"]
     assert np.array_equal(np.unique(fancy_test_array[columns]), fancy_test_array[columns])
 
 
-def test_getitem_array_slice(fancy_test_array):
+def test_getitem_array_slice(fancy_test_array: FancyTestArray):
     assert fancy_test_array.data[0:2].tolist() == fancy_test_array[0:2].tolist()
 
 
-def test_getitem_with_array_mask(fancy_test_array):
+def test_getitem_with_array_mask(fancy_test_array: FancyTestArray):
     mask = np.array([True, False, True])
     assert isinstance(fancy_test_array[mask], FancyArray)
     assert np.array_equal(fancy_test_array.data[mask], fancy_test_array[mask].data)
 
 
-def test_getitem_with_tuple_mask(fancy_test_array):
+def test_getitem_with_tuple_mask(fancy_test_array: FancyTestArray):
     mask = (True, False, True)
     assert isinstance(fancy_test_array[mask], FancyArray)
     assert np.array_equal(fancy_test_array.data[mask], fancy_test_array[mask].data)
 
 
-def test_getitem_with_list_mask(fancy_test_array):
+def test_getitem_with_list_mask(fancy_test_array: FancyTestArray):
     mask = [True, False, True]
     assert isinstance(fancy_test_array[mask], FancyArray)
     assert np.array_equal(fancy_test_array.data[mask], fancy_test_array[mask].data)
@@ -102,54 +102,54 @@ def test_getitem_with_empty_list_mask():
     assert np.array_equal(array.data[mask], array[mask].data)
 
 
-def test_setitem_with_index(fancy_test_array):
+def test_setitem_with_index(fancy_test_array: FancyTestArray):
     fancy_test_array[0] = (9, 9, 9, 9, 9)
     assert [9, 2, 3] == fancy_test_array.id.tolist()
 
 
-def test_setitem_with_mask(fancy_test_array):
+def test_setitem_with_mask(fancy_test_array: FancyTestArray):
     mask = np.array([True, False, True])
     fancy_test_array[mask] = (9, 9, 9, 9, 9)
     assert [9, 2, 9] == fancy_test_array.id.tolist()
 
 
-def test_setitem_as_fancy_array_with_mask(fancy_test_array):
+def test_setitem_as_fancy_array_with_mask(fancy_test_array: FancyTestArray):
     mask = np.array([True, False, True])
     fancy_test_array[mask] = FancyTestArray.zeros(2)
     assert_array_equal([EMPTY_ID, 2, EMPTY_ID], fancy_test_array.id)
 
 
-def test_setitem_as_fancy_array_with_mask_too_large(fancy_test_array):
+def test_setitem_as_fancy_array_with_mask_too_large(fancy_test_array: FancyTestArray):
     mask = np.array([True, False, True])
     with pytest.raises(ValueError):
         fancy_test_array[mask] = FancyTestArray.zeros(3)
 
 
-def test_set_non_existing_field(fancy_test_array):
+def test_set_non_existing_field(fancy_test_array: FancyTestArray):
     with pytest.raises(AttributeError):
         fancy_test_array.non_existing_field = 123
 
 
-def test_set_callable(fancy_test_array):
+def test_set_callable(fancy_test_array: FancyTestArray):
     with pytest.raises(AttributeError):
         fancy_test_array.filter = 123
 
 
-def test_contains(fancy_test_array):
+def test_contains(fancy_test_array: FancyTestArray):
     assert fancy_test_array[0] in fancy_test_array
 
 
-def test_non_existing_method(fancy_test_array):
+def test_non_existing_method(fancy_test_array: FancyTestArray):
     with pytest.raises(AttributeError):
         # pylint: disable=no-member
         fancy_test_array.non_existing_method()
 
 
-def test_array_equal(fancy_test_array):
+def test_array_equal(fancy_test_array: FancyTestArray):
     assert fp.array_equal(fancy_test_array, fancy_test_array.copy())
 
 
-def test_array_not_equal(fancy_test_array):
+def test_array_not_equal(fancy_test_array: FancyTestArray):
     different_array = fancy_test_array.copy()
     different_array.test_int = 99
     assert not fp.array_equal(fancy_test_array, different_array)
@@ -203,60 +203,60 @@ def test_is_empty_bool():
     assert_array_equal(array.is_empty("test_bool"), [True, False])
 
 
-def test_unique(fancy_test_array):
+def test_unique(fancy_test_array: FancyTestArray):
     duplicate_array = fp.concatenate(fancy_test_array, fancy_test_array)
     unique_array = fp.unique(duplicate_array)
     assert fp.array_equal(unique_array, fancy_test_array)
 
 
-def test_unique_with_nan_values(fancy_test_array):
+def test_unique_with_nan_values(fancy_test_array: FancyTestArray):
     fancy_test_array.test_float = np.nan
     duplicate_array = fp.concatenate(fancy_test_array, fancy_test_array)
     with pytest.raises(NotImplementedError):
         fp.unique(duplicate_array)
 
 
-def test_unique_return_inverse(fancy_test_array):
+def test_unique_return_inverse(fancy_test_array: FancyTestArray):
     duplicate_array = fp.concatenate(fancy_test_array, fancy_test_array)
     unique_array, inverse = fp.unique(duplicate_array, return_inverse=True)
     assert fp.array_equal(unique_array, fancy_test_array)
     assert_array_equal(inverse, [0, 1, 2, 0, 1, 2])
 
 
-def test_unique_return_counts_and_inverse(fancy_test_array):
+def test_unique_return_counts_and_inverse(fancy_test_array: FancyTestArray):
     duplicate_array = fp.concatenate(fancy_test_array, fancy_test_array)
     unique_array, inverse, counts = fp.unique(duplicate_array, return_counts=True, return_inverse=True)
     assert fp.array_equal(unique_array, fancy_test_array)
     assert_array_equal(counts, [2, 2, 2])
     assert_array_equal(inverse, [0, 1, 2, 0, 1, 2])
 
 
-def test_sort(fancy_test_array):
+def test_sort(fancy_test_array: FancyTestArray):
     assert_array_equal(fancy_test_array.test_float, [4.0, 4.0, 1.0])
     fancy_test_array.sort(order="test_float")
     assert_array_equal(fancy_test_array.test_float, [1.0, 4.0, 4.0])
 
 
-def test_copy_function(fancy_test_array):
+def test_copy_function(fancy_test_array: FancyTestArray):
     array_copy = copy(fancy_test_array)
     array_copy.test_int = 123
     assert not id(fancy_test_array) == id(array_copy)
     assert not fancy_test_array.test_int[0] == array_copy.test_int[0]
 
 
-def test_copy_method(fancy_test_array):
+def test_copy_method(fancy_test_array: FancyTestArray):
     array_copy = fancy_test_array.copy()
     array_copy.test_int = 123
     assert not id(fancy_test_array.data) == id(array_copy.data)
-    assert not fancy_test_array.test_int[0] == array_copy.test_int[0]
+    assert not fancy_test_array.test_int[0] == array_copy.test_int[0]  # type: ignore
 
 
-def test_prevent_np_unique_on_fancy_array(fancy_test_array):
+def test_prevent_np_unique_on_fancy_array(fancy_test_array: FancyTestArray):
     with pytest.raises(TypeError):
         np.unique(fancy_test_array)
 
 
-def test_prevent_np_sort_on_fancy_array(fancy_test_array):
+def test_prevent_np_sort_on_fancy_array(fancy_test_array: FancyTestArray):
     with pytest.raises(TypeError):
         np.sort(fancy_test_array)
 
@@ -278,7 +278,7 @@ def test_string_inherit_string_length():
     assert_array_equal(array.extra_string, ["b" * 100])
 
 
-def test_shuffle_array(fancy_test_array):
+def test_shuffle_array(fancy_test_array: FancyTestArray):
     rng = np.random.default_rng(0)
     rng.shuffle(fancy_test_array.data)
     assert_array_equal(fancy_test_array.id, [3, 1, 2])