From a20d6785373375e55cb62121bad87d7835382829 Mon Sep 17 00:00:00 2001 From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> Date: Fri, 31 Oct 2025 11:52:49 +0100 Subject: [PATCH 1/4] Make __getitem__ typesafe by restricting input Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> --- .../_core/model/arrays/base/array.py | 37 ++++++++++-------- .../_core/model/arrays/pgm_arrays.py | 8 ++-- tests/unit/model/arrays/test_array.py | 39 +++++++++++-------- tests/unit/model/arrays/test_modify.py | 4 +- 4 files changed, 49 insertions(+), 39 deletions(-) diff --git a/src/power_grid_model_ds/_core/model/arrays/base/array.py b/src/power_grid_model_ds/_core/model/arrays/base/array.py index c4bc09f..55b516a 100644 --- a/src/power_grid_model_ds/_core/model/arrays/base/array.py +++ b/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)) diff --git a/src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py b/src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py index 9280d72..ee64946 100644 --- a/src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py +++ b/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): diff --git a/tests/unit/model/arrays/test_array.py b/tests/unit/model/arrays/test_array.py index 8e4b7c5..1e6ac91 100644 --- a/tests/unit/model/arrays/test_array.py +++ b/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]) diff --git a/tests/unit/model/arrays/test_modify.py b/tests/unit/model/arrays/test_modify.py index 598df90..e690bc9 100644 --- a/tests/unit/model/arrays/test_modify.py +++ b/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 From ab3dd19df27ac58936fa1cdfbe41bc7b6da02b7b Mon Sep 17 00:00:00 2001 From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> Date: Fri, 31 Oct 2025 12:36:40 +0100 Subject: [PATCH 2/4] revert some changes Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> --- .../_core/model/arrays/base/array.py | 24 ++++++++---- .../_core/model/arrays/pgm_arrays.py | 8 ++-- tests/unit/model/arrays/test_array.py | 37 +++++++++++-------- 3 files changed, 42 insertions(+), 27 deletions(-) diff --git a/src/power_grid_model_ds/_core/model/arrays/base/array.py b/src/power_grid_model_ds/_core/model/arrays/base/array.py index 55b516a..5d6810e 100644 --- a/src/power_grid_model_ds/_core/model/arrays/base/array.py +++ b/src/power_grid_model_ds/_core/model/arrays/base/array.py @@ -153,10 +153,12 @@ 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): @@ -164,13 +166,19 @@ def __getitem__(self, 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)) diff --git a/src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py b/src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py index ee64946..9280d72 100644 --- a/src/power_grid_model_ds/_core/model/arrays/pgm_arrays.py +++ b/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): diff --git a/tests/unit/model/arrays/test_array.py b/tests/unit/model/arrays/test_array.py index 1e6ac91..3875b99 100644 --- a/tests/unit/model/arrays/test_array.py +++ b/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]) From f2aabd6cd515f0454c593aabc0ab4a12ab7d8415 Mon Sep 17 00:00:00 2001 From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> Date: Fri, 31 Oct 2025 12:38:34 +0100 Subject: [PATCH 3/4] revert some changes Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> --- tests/unit/model/arrays/test_modify.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/model/arrays/test_modify.py b/tests/unit/model/arrays/test_modify.py index e690bc9..598df90 100644 --- a/tests/unit/model/arrays/test_modify.py +++ b/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.data[["test_str", "test_int"]] + sub_array = fancy_test_array[["test_str", "test_int"]] different_array = FancyNonIdArray.zeros(10) - different_sub_array = different_array.data[["test_str", "test_int"]] + different_sub_array = different_array[["test_str", "test_int"]] concatenated = np.concatenate([sub_array, different_sub_array]) assert concatenated.size == 13 From 65434f0d6cdb4e998048e777903627c4829b353e Mon Sep 17 00:00:00 2001 From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> Date: Wed, 12 Nov 2025 13:57:27 +0100 Subject: [PATCH 4/4] merge main Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com> --- .github/workflows/ci.yml | 6 ++-- .../_core/model/arrays/base/array.py | 32 ++++++++++--------- tests/unit/model/arrays/test_array.py | 1 + 3 files changed, 21 insertions(+), 18 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1beb8f2..d65099a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -62,7 +62,7 @@ jobs: uv build - name: Store built wheel file - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v5 with: name: power-grid-model-ds path: dist/ @@ -91,7 +91,7 @@ jobs: run: uv tool install poethepoet - name: Load built wheel file - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v6 with: name: power-grid-model-ds path: dist/ @@ -117,7 +117,7 @@ jobs: uses: actions/checkout@v5 - name: Load built wheel file - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v6 with: name: power-grid-model-ds path: dist/ diff --git a/src/power_grid_model_ds/_core/model/arrays/base/array.py b/src/power_grid_model_ds/_core/model/arrays/base/array.py index 5d6810e..b62aa74 100644 --- a/src/power_grid_model_ds/_core/model/arrays/base/array.py +++ b/src/power_grid_model_ds/_core/model/arrays/base/array.py @@ -59,14 +59,14 @@ class FancyArray(ABC): _defaults: dict[str, Any] = {} _str_lengths: dict[str, int] = {} - def __init__(self: Self, *args, data: NDArray | None = None, **kwargs): + def __init__(self, *args, data: NDArray | None = None, **kwargs): if data is None: self._data = build_array(*args, dtype=self.get_dtype(), defaults=self.get_defaults(), **kwargs) else: self._data = data @property - def data(self: Self) -> NDArray: + def data(self) -> NDArray: return self._data @classmethod @@ -110,7 +110,7 @@ def get_dtype(cls): dtype_list.append((name, dtype)) return np.dtype(dtype_list) - def __repr__(self: Self) -> str: + def __repr__(self) -> str: try: data = getattr(self, "data") if data.size > 3: @@ -125,7 +125,7 @@ def __str__(self) -> str: def __len__(self) -> int: return len(self._data) - def __iter__(self: Self): + def __iter__(self): for record in self._data: yield self.__class__(data=np.array([record])) @@ -190,16 +190,18 @@ def __contains__(self: Self, item: Self) -> bool: return item.data in self._data return False - def __hash__(self: Self): + def __hash__(self): return hash(f"{self.__class__} {self}") - def __eq__(self: Self, other): - return self._data.__eq__(other.data) + def __eq__(self, other): + if not isinstance(other, self.__class__): + return False + return self.data.__eq__(other.data) - def __copy__(self: Self): + def __copy__(self): return self.__class__(data=copy(self._data)) - def copy(self: Self): + def copy(self): """Return a copy of this array including its data""" return copy(self) @@ -294,7 +296,7 @@ def get( return self.__class__(data=apply_get(*args, array=self._data, mode_=mode_, **kwargs)) def filter_mask( - self: Self, + self, *args: int | Iterable[int] | np.ndarray, mode_: Literal["AND", "OR"] = "AND", **kwargs: Any | list[Any] | np.ndarray, @@ -302,7 +304,7 @@ def filter_mask( return get_filter_mask(*args, array=self._data, mode_=mode_, **kwargs) def exclude_mask( - self: Self, + self, *args: int | Iterable[int] | np.ndarray, mode_: Literal["AND", "OR"] = "AND", **kwargs: Any | list[Any] | np.ndarray, @@ -312,7 +314,7 @@ def exclude_mask( def re_order(self: Self, new_order: ArrayLike, column: str = "id") -> Self: return self.__class__(data=re_order(self._data, new_order, column=column)) - def update_by_id(self: Self, ids: ArrayLike, allow_missing: bool = False, **kwargs) -> None: + def update_by_id(self, ids: ArrayLike, allow_missing: bool = False, **kwargs) -> None: try: _ = update_by_id(self._data, ids, allow_missing, **kwargs) except ValueError as error: @@ -325,13 +327,13 @@ def get_updated_by_id(self: Self, ids: ArrayLike, allow_missing: bool = False, * except ValueError as error: raise ValueError(f"Cannot update {self.__class__.__name__}. {error}") from error - def check_ids(self: Self, return_duplicates: bool = False) -> NDArray | None: + def check_ids(self, return_duplicates: bool = False) -> NDArray | None: return check_ids(self._data, return_duplicates=return_duplicates) - def as_table(self: Self, column_width: int | str = "auto", rows: int = 10) -> str: + def as_table(self, column_width: int | str = "auto", rows: int = 10) -> str: return convert_array_to_string(self, column_width=column_width, rows=rows) - def as_df(self: Self): + def as_df(self): """Convert to pandas DataFrame""" if pandas is None: raise ImportError("pandas is not installed") diff --git a/tests/unit/model/arrays/test_array.py b/tests/unit/model/arrays/test_array.py index 3875b99..c727687 100644 --- a/tests/unit/model/arrays/test_array.py +++ b/tests/unit/model/arrays/test_array.py @@ -65,6 +65,7 @@ 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"))