Make __getitem__ typesafe by restricting input

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

Author: Thijss

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

@@ -6,7 +6,7 @@
 from collections import namedtuple
 from copy import copy
 from functools import lru_cache
-from typing import Any, Iterable, Literal, Type, TypeVar
+from typing import Any, Iterable, Literal, Type, TypeVar, overload
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
@@ -152,20 +152,25 @@ def __setattr__(self: Self, attr: str, value: object) -> None:
         except (AttributeError, ValueError) as error:
             raise AttributeError(f"Cannot set attribute {attr} on {self.__class__.__name__}") from error
 
-    def __getitem__(self: Self, item):
-        """Used by for-loops, slicing [0:3], column-access ['id'], row-access [0], multi-column access.
-        Note: If a single item is requested, return a named tuple instead of a np.void object.
-        """
-
-        result = self._data.__getitem__(item)
-
-        if isinstance(item, (list, tuple)) and (len(item) == 0 or np.array(item).dtype.type is np.bool_):
-            return self.__class__(data=result)
-        if isinstance(item, (str, list, tuple)):
-            return result
-        if isinstance(result, np.void):
-            return self.__class__(data=np.array([result]))
-        return self.__class__(data=result)
+    @overload
+    def __getitem__(self: Self, item: slice | int | NDArray[np.bool_]) -> Self: ...
+
+    @overload
+    def __getitem__(self, item: str | NDArray[np.str_]) -> NDArray[Any]: ...
+
+    def __getitem__(self, item):
+        if isinstance(item, slice | int):
+            new_data = self._data[item]
+            if new_data.shape == ():
+                new_data = np.array([new_data])
+            return self.__class__(data=new_data)
+        if isinstance(item, np.ndarray) and item.dtype == np.bool_:
+            return self.__class__(data=self._data[item])
+        if isinstance(item, np.ndarray) and item.size == 0:
+            return self.__class__(data=self._data[[]])
+        if isinstance(item, str):
+            return self._data[item]
+        raise NotImplementedError(f"FancyArray[{type(item)}] is not supported. Use FancyArray.data instead.")
 
     def __setitem__(self: Self, key, value):
         if isinstance(value, FancyArray):
@@ -330,4 +335,4 @@ def from_extended(cls: Type[Self], extended: Self) -> Self:
         if not isinstance(extended, cls):
             raise TypeError(f"Extended array must be of type {cls.__name__}, got {type(extended).__name__}")
         dtype = cls.get_dtype()
-        return cls(data=np.array(extended[list(dtype.names)], dtype=dtype))
+        return cls(data=np.array(extended.data[list(dtype.names)], dtype=dtype))

src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py

@@ -71,7 +71,7 @@ def filter_parallel(self, n_parallel: int, mode: Literal["eq", "neq"]) -> "Branc
             - when n_parallel is 1 and mode is 'eq', the function returns branches that are not parallel.
             - when n_parallel is 1 and mode is 'neq', the function returns branches that are parallel.
         """
-        _, index, counts = np.unique(self[["from_node", "to_node"]], return_counts=True, return_index=True)
+        _, index, counts = np.unique(self.data[["from_node", "to_node"]], return_counts=True, return_index=True)
 
         match mode:
             case "eq":
@@ -82,9 +82,9 @@ def filter_parallel(self, n_parallel: int, mode: Literal["eq", "neq"]) -> "Branc
                 raise ValueError(f"mode {mode} not supported")
 
         if mode == "eq" and n_parallel == 1:
-            return self[index][counts_mask]
-        filtered_branches = self[index][counts_mask]
-        return self.filter(from_node=filtered_branches.from_node, to_node=filtered_branches.to_node)
+            return self.__class__(self.data[index][counts_mask])
+        filtered_branches = self.data[index][counts_mask]
+        return self.filter(from_node=filtered_branches["from_node"], to_node=filtered_branches["to_node"])
 
 
 class LinkArray(Link, BranchArray):

tests/unit/model/arrays/test_array.py

@@ -65,17 +65,22 @@ def test_getitem_array_one_column(fancy_test_array: FancyTestArray):
 
 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"))
+    assert_array_equal(fancy_test_array.data[columns].dtype.names, ("id", "test_int", "test_float"))
 
 
-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_index(fancy_test_array: FancyTestArray):
+    assert fancy_test_array[0].data.tolist() == fancy_test_array.data[0:1].tolist()
+
+
+def test_getitem_array_nested_index(fancy_test_array: FancyTestArray):
+    nested_array = fancy_test_array[0][0][0][0][0][0]
+    assert isinstance(nested_array, FancyArray)
+    assert nested_array.data.shape == (1,)
+    assert nested_array.data.tolist() == fancy_test_array.data[0:1].tolist()
 
 
 def test_getitem_array_slice(fancy_test_array: FancyTestArray):
-    assert fancy_test_array.data[0:2].tolist() == fancy_test_array[0:2].tolist()
+    assert fancy_test_array[0:2].data.tolist() == fancy_test_array.data[0:2].tolist()
 
 
 def test_getitem_with_array_mask(fancy_test_array: FancyTestArray):
@@ -86,19 +91,19 @@ def test_getitem_with_array_mask(fancy_test_array: FancyTestArray):
 
 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)
+    with pytest.raises(NotImplementedError):
+        fancy_test_array[mask]  # type: ignore[call-overload]  # noqa
 
 
 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)
+    with pytest.raises(NotImplementedError):
+        fancy_test_array[mask]  # type: ignore[call-overload]  # noqa
 
 
 def test_getitem_with_empty_list_mask():
     array = FancyTestArray()
-    mask = []
+    mask = np.array([], dtype=bool)
     assert isinstance(array[mask], FancyArray)
     assert np.array_equal(array.data[mask], array[mask].data)
 
@@ -233,16 +238,16 @@ def test_unique_return_counts_and_inverse(fancy_test_array: FancyTestArray):
 
 
 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])
+    assert_array_equal(fancy_test_array["test_float"], [4.0, 4.0, 1.0])
+    fancy_test_array.data.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: FancyTestArray):
     array_copy = copy(fancy_test_array)
-    array_copy.test_int = 123
+    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]
+    assert not fancy_test_array["test_int"][0] == array_copy["test_int"][0]
 
 
 def test_copy_method(fancy_test_array: FancyTestArray):
@@ -302,4 +307,4 @@ def test_from_extended_array():
 
     array = LineArray.from_extended(extended_array)
     assert not isinstance(array, ExtendedLineArray)
-    array_equal_with_nan(array.data, extended_array[array.columns])
+    array_equal_with_nan(array.data, extended_array.data[array.columns])

tests/unit/model/arrays/test_modify.py

@@ -128,10 +128,10 @@ def test_concatenate_different_ndarray(self, fancy_test_array: FancyTestArray):
             fp.concatenate(fancy_test_array, different_array.data)
 
     def test_concatenate_different_fancy_array_same_dtype(self, fancy_test_array: FancyTestArray):
-        sub_array = fancy_test_array[["test_str", "test_int"]]
+        sub_array = fancy_test_array.data[["test_str", "test_int"]]
 
         different_array = FancyNonIdArray.zeros(10)
-        different_sub_array = different_array[["test_str", "test_int"]]
+        different_sub_array = different_array.data[["test_str", "test_int"]]
 
         concatenated = np.concatenate([sub_array, different_sub_array])
         assert concatenated.size == 13