switch extended arrays

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

Author: Thijss

tests/fixtures/arrays.py

@@ -5,6 +5,7 @@
 import numpy as np
 from numpy._typing import NDArray
 
+from power_grid_model_ds._core.model.arrays import LineArray, NodeArray
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 from power_grid_model_ds._core.model.dtypes.sensors import NDArray3
 
@@ -36,15 +37,6 @@ class DifferentFancyTestArray(FancyArray):
     test_bool2: NDArray[np.bool_]
 
 
-class ExtendedFancyTestArray(FancyTestArray):
-    """Test array with some extended attributes"""
-
-    test_int2: NDArray[np.int64]
-    test_float2: NDArray[np.float64]
-    test_str2: NDArray[np.str_]
-    test_bool2: NDArray[np.bool_]
-
-
 class LongColumnNameFancyTestArray(FancyArray):
     """Test array with a very long column name"""
 
@@ -66,3 +58,19 @@ class FancyTestArray3(FancyArray):
 
     test_float1: NDArray3[np.float64]
     test_float2: NDArray3[np.float64]
+
+
+class ExtendedNodeArray(NodeArray):
+    """Extends the node array with an output value"""
+
+    _defaults = {"u": 0}
+
+    u: NDArray[np.float64]
+
+
+class ExtendedLineArray(LineArray):
+    """Extends the line array with an output value"""
+
+    _defaults = {"i_from": 0}
+
+    i_from: NDArray[np.float64]

tests/integration/loadflow/test_power_grid_model.py

@@ -23,27 +23,12 @@
 )
 from power_grid_model_ds._core.model.arrays.pgm_arrays import TransformerTapRegulatorArray
 from power_grid_model_ds._core.model.grids.base import Grid
+from tests.fixtures.arrays import ExtendedLineArray, ExtendedNodeArray
 from tests.unit.model.grids.test_custom_grid import CustomGrid
 
 # pylint: disable=missing-function-docstring,missing-class-docstring
 
 
-class ExtendedNodeArray(NodeArray):
-    """Extends the node array with an output value"""
-
-    _defaults = {"u": 0}
-
-    u: NDArray[np.float64]
-
-
-class ExtendedLineArray(LineArray):
-    """Extends the line array with an output value"""
-
-    _defaults = {"i_from": 0}
-
-    i_from: NDArray[np.float64]
-
-
 def test_load_flow_on_random():
     """Tests the power flow on a randomly configured grid"""
     grid_generator = RadialGridGenerator(grid_class=Grid, nr_nodes=5, nr_sources=1, nr_nops=0)

tests/unit/model/arrays/test_array.py

@@ -11,10 +11,10 @@
 
 from power_grid_model_ds._core import fancypy as fp
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
-from power_grid_model_ds._core.model.arrays.pgm_arrays import TransformerArray
+from power_grid_model_ds._core.model.arrays.pgm_arrays import LineArray, TransformerArray
 from power_grid_model_ds._core.model.constants import EMPTY_ID, empty
 from tests.conftest import FancyTestArray
-from tests.fixtures.arrays import ExtendedFancyTestArray, FancyTestArray3
+from tests.fixtures.arrays import ExtendedLineArray, FancyTestArray3
 
 # pylint: disable=missing-function-docstring
 
@@ -292,7 +292,7 @@ def test_overflow_value():
 
 
 def test_from_extended_array():
-    extended_array = ExtendedFancyTestArray.zeros(2)
-    array = FancyTestArray.from_extended(extended_array)
-    assert not isinstance(array, ExtendedFancyTestArray)
+    extended_array = ExtendedLineArray.zeros(2)
+    array = LineArray.from_extended(extended_array)
+    assert not isinstance(array, ExtendedLineArray)
     assert_array_equal(array.data, extended_array[array.columns])