revert some changes

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

Author: Thijss

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

@@ -153,24 +153,32 @@ def __setattr__(self: Self, attr: str, value: object) -> None:
             raise AttributeError(f"Cannot set attribute {attr} on {self.__class__.__name__}") from error
 
     @overload
-    def __getitem__(self: Self, item: slice | int | NDArray[np.bool_]) -> Self: ...
+    def __getitem__(
+        self: Self, item: slice | int | NDArray[np.bool_] | list[bool] | NDArray[np.int_] | list[int]
+    ) -> Self: ...
 
     @overload
-    def __getitem__(self, item: str | NDArray[np.str_]) -> NDArray[Any]: ...
+    def __getitem__(self, item: str | NDArray[np.str_] | list[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.")
+        if (isinstance(item, np.ndarray) and item.size == 0) or (isinstance(item, list | tuple) and len(item) == 0):
+            return self.__class__(data=self._data[[]])
+        if isinstance(item, list | np.ndarray):
+            item_array = np.array(item)
+            if item_array.dtype == np.bool_ or np.issubdtype(item_array.dtype, np.int_):
+                return self.__class__(data=self._data[item_array])
+            if np.issubdtype(item_array.dtype, np.str_):
+                return self._data[item_array.tolist()]
+        raise NotImplementedError(
+            f"FancyArray[{type(item).__name__}] is not supported. Try FancyArray.data[{type(item).__name__}] instead."
+        )
 
     def __setitem__(self: Self, key, value):
         if isinstance(value, FancyArray):
@@ -335,4 +343,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.data[list(dtype.names)], dtype=dtype))
+        return cls(data=np.array(extended[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.data[["from_node", "to_node"]], return_counts=True, return_index=True)
+        _, index, counts = np.unique(self[["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.__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"])
+            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)
 
 
 class LinkArray(Link, BranchArray):

tests/unit/model/arrays/test_array.py

@@ -65,7 +65,12 @@ 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_array_equal(fancy_test_array.data[columns].dtype.names, ("id", "test_int", "test_float"))
+    assert_array_equal(fancy_test_array[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):
@@ -80,7 +85,7 @@ def test_getitem_array_nested_index(fancy_test_array: FancyTestArray):
 
 
 def test_getitem_array_slice(fancy_test_array: FancyTestArray):
-    assert fancy_test_array[0:2].data.tolist() == fancy_test_array.data[0:2].tolist()
+    assert fancy_test_array.data[0:2].tolist() == fancy_test_array[0:2].tolist()
 
 
 def test_getitem_with_array_mask(fancy_test_array: FancyTestArray):
@@ -89,21 +94,23 @@ def test_getitem_with_array_mask(fancy_test_array: FancyTestArray):
     assert np.array_equal(fancy_test_array.data[mask], fancy_test_array[mask].data)
 
 
-def test_getitem_with_tuple_mask(fancy_test_array: FancyTestArray):
-    mask = (True, False, True)
-    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)
+
+
+def test_getitem_with_tuple_mask(fancy_test_array: FancyTestArray):
+    # Numpy gives unexpected results with tuple masks. Therefore, we raise NotImplementedError here.
+    # e.g: np.array([1,2,3])[(True, False, True)] returns an empty array (array([], shape=(0, 3), dtype=int64)
+    mask = (True, False, True)
     with pytest.raises(NotImplementedError):
         fancy_test_array[mask]  # type: ignore[call-overload]  # noqa
 
 
 def test_getitem_with_empty_list_mask():
     array = FancyTestArray()
-    mask = np.array([], dtype=bool)
+    mask = []
     assert isinstance(array[mask], FancyArray)
     assert np.array_equal(array.data[mask], array[mask].data)
 
@@ -238,16 +245,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.data.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.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):
@@ -307,4 +314,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.data[array.columns])
+    array_equal_with_nan(array.data, extended_array[array.columns])