Add FancyArray.from_extended

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

Author: Thijss

src/power_grid_model_ds/_core/model/arrays/base/array.py

@@ -323,3 +323,13 @@ def as_df(self: Self):
         if pandas is None:
             raise ImportError("pandas is not installed")
         return pandas.DataFrame(self._data)
+
+    @classmethod
+    def from_extended(cls: Type[Self], extended: Self) -> Self:
+        """Create an instance from an extended array."""
+        if not isinstance(extended, cls):
+            raise TypeError(f"Extended array must be of type {cls.__name__}, got {type(extended).__name__}")
+        new_array = cls.zeros(len(extended), empty_id=False)
+        for column in cls.get_dtype().names:
+            new_array[column] = extended[column]
+        return new_array

tests/fixtures/arrays.py

@@ -36,6 +36,15 @@ 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"""
 

tests/unit/model/arrays/test_array.py

@@ -14,7 +14,7 @@
 from power_grid_model_ds._core.model.arrays.pgm_arrays import TransformerArray
 from power_grid_model_ds._core.model.constants import EMPTY_ID, empty
 from tests.conftest import FancyTestArray
-from tests.fixtures.arrays import FancyTestArray3
+from tests.fixtures.arrays import ExtendedFancyTestArray, FancyTestArray3
 
 # pylint: disable=missing-function-docstring
 
@@ -289,3 +289,10 @@ def test_overflow_value():
     with pytest.raises(OverflowError):
         transformer.tap_min = -167
     assert transformer.tap_min == -128
+
+
+def test_from_extended_array():
+    extended_array = ExtendedFancyTestArray.zeros(2)
+    array = FancyTestArray.from_extended(extended_array)
+    assert not isinstance(array, ExtendedFancyTestArray)
+    assert_array_equal(array.data, extended_array[array.columns])