From 2fe88dbac3b5216e98f69fb134320e0c7cebe73d Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Tue, 31 Oct 2023 20:47:47 +0100 Subject: [PATCH 1/7] Update _typing.py --- xarray/namedarray/_typing.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 0b972e19539..751a5e47cba 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -19,7 +19,7 @@ _T = TypeVar("_T") _T_co = TypeVar("_T_co", covariant=True) - +_dtype = np.dtype _DType = TypeVar("_DType", bound=np.dtype[Any]) _DType_co = TypeVar("_DType_co", covariant=True, bound=np.dtype[Any]) # A subset of `npt.DTypeLike` that can be parametrized w.r.t. `np.generic` @@ -55,9 +55,15 @@ def dtype(self) -> _DType_co: _Dims = tuple[_Dim, ...] _DimsLike = Union[str, Iterable[_Dim]] -_AttrsLike = Union[Mapping[Any, Any], None] -_dtype = np.dtype +# https://data-apis.org/array-api/latest/API_specification/indexing.html +# TODO: np.array_api doesn't allow None for some reason, maybe they're +# recommending to use expand_dims? +_IndexKey = Union[int, slice, "ellipsis"] +_IndexKeys = tuple[Union[_IndexKey, None], ...] +_IndexKeyLike = Union[_IndexKey, _IndexKeys] + +_AttrsLike = Union[Mapping[Any, Any], None] class _SupportsReal(Protocol[_T_co]): From b483a186e605cbc8fd678f0dcbfb18e08467c502 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Tue, 31 Oct 2023 20:47:54 +0100 Subject: [PATCH 2/7] Update _typing.py --- xarray/namedarray/_typing.py | 38 ++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 751a5e47cba..ba9a3c3ace5 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -105,6 +105,23 @@ class _arrayfunction( Corresponds to np.ndarray. """ + @overload + def __getitem__(self, key: _IndexKeyLike) -> Any: + ... + + @overload + def __getitem__( + self, key: _arrayfunction[Any, Any] | tuple[_arrayfunction[Any, Any], ...] + ) -> _arrayfunction[Any, _DType_co]: + ... + + def __getitem__( + self, + key: _IndexKeyLike | _arrayfunction[Any, Any], + /, + ) -> _arrayfunction[Any, _DType_co] | Any: + ... + @overload def __array__(self, dtype: None = ..., /) -> np.ndarray[Any, _DType_co]: ... @@ -157,6 +174,27 @@ class _arrayapi(_array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType Corresponds to np.ndarray. """ + # TODO: Only integer _arrayapi: + # def __getitem__( + # self, + # key: Union[ + # int, + # slice, + # "ellipsis", + # tuple[Union[int, slice, "ellipsis", None], ...], + # _arrayapi[Any, Any], + # ], + # /, + # ) -> _arrayapi[Any, _DType_co]: + # ... + + def __getitem__( + self, + key: Any, + /, + ) -> _arrayapi[Any, _DType_co]: + ... + def __array_namespace__(self) -> ModuleType: ... From 1b193bcf4ad31ebecfdaa5f891dd3ef3bb12bf63 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Tue, 31 Oct 2023 20:50:53 +0100 Subject: [PATCH 3/7] Update test_namedarray.py --- xarray/tests/test_namedarray.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index 448e8cf819a..33bfe352ddc 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -22,6 +22,7 @@ _AttrsLike, _DimsLike, _DType, + _IndexKeyLike, _Shape, duckarray, ) @@ -53,6 +54,11 @@ class CustomArrayIndexable( ExplicitlyIndexed, Generic[_ShapeType_co, _DType_co], ): + def __getitem__( + self, key: _IndexKeyLike | CustomArrayIndexable[Any, Any], / + ) -> CustomArrayIndexable[Any, _DType_co]: + return type(self)(array=self.array[key]) + def __array_namespace__(self) -> ModuleType: return np From 451cb4f959a9e5c918bffd1335f3f47371f8c692 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Fri, 3 Nov 2023 00:01:56 +0100 Subject: [PATCH 4/7] fixes --- xarray/namedarray/_typing.py | 32 ++++++++++---------------------- xarray/tests/test_namedarray.py | 11 ++++++++++- 2 files changed, 20 insertions(+), 23 deletions(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index ba9a3c3ace5..70f44b39dd2 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -60,7 +60,7 @@ def dtype(self) -> _DType_co: # TODO: np.array_api doesn't allow None for some reason, maybe they're # recommending to use expand_dims? _IndexKey = Union[int, slice, "ellipsis"] -_IndexKeys = tuple[Union[_IndexKey, None], ...] +_IndexKeys = tuple[Union[_IndexKey], ...] _IndexKeyLike = Union[_IndexKey, _IndexKeys] _AttrsLike = Union[Mapping[Any, Any], None] @@ -106,18 +106,20 @@ class _arrayfunction( """ @overload - def __getitem__(self, key: _IndexKeyLike) -> Any: + def __getitem__( + self, key: _arrayfunction[Any, Any] | tuple[_arrayfunction[Any, Any], ...], / + ) -> _arrayfunction[Any, _DType_co]: ... @overload - def __getitem__( - self, key: _arrayfunction[Any, Any] | tuple[_arrayfunction[Any, Any], ...] - ) -> _arrayfunction[Any, _DType_co]: + def __getitem__(self, key: _IndexKeyLike, /) -> Any: ... def __getitem__( self, - key: _IndexKeyLike | _arrayfunction[Any, Any], + key: _IndexKeyLike + | _arrayfunction[Any, Any] + | tuple[_arrayfunction[Any, Any], ...], /, ) -> _arrayfunction[Any, _DType_co] | Any: ... @@ -174,25 +176,11 @@ class _arrayapi(_array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType Corresponds to np.ndarray. """ - # TODO: Only integer _arrayapi: - # def __getitem__( - # self, - # key: Union[ - # int, - # slice, - # "ellipsis", - # tuple[Union[int, slice, "ellipsis", None], ...], - # _arrayapi[Any, Any], - # ], - # /, - # ) -> _arrayapi[Any, _DType_co]: - # ... - def __getitem__( self, - key: Any, + key: _IndexKeyLike | Any, # TODO: Any should be _arrayapi /, - ) -> _arrayapi[Any, _DType_co]: + ) -> _arrayapi[Any, Any]: ... def __array_namespace__(self) -> ModuleType: diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index 33bfe352ddc..b3ae5781d99 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -19,6 +19,7 @@ from numpy.typing import ArrayLike, DTypeLike, NDArray from xarray.namedarray._typing import ( + _arrayapi, _AttrsLike, _DimsLike, _DType, @@ -57,7 +58,15 @@ class CustomArrayIndexable( def __getitem__( self, key: _IndexKeyLike | CustomArrayIndexable[Any, Any], / ) -> CustomArrayIndexable[Any, _DType_co]: - return type(self)(array=self.array[key]) + if isinstance(key, CustomArrayIndexable): + if isinstance(key.array, type(self.array)): + # TODO: key.array is duckarray here, can it be narrowed down further? + # an _arrayapi cannot be used on a _arrayfunction for example. + return type(self)(array=self.array[key.array]) # type: ignore[index] + else: + raise TypeError("key must have the same array type as self") + else: + return type(self)(array=self.array[key]) def __array_namespace__(self) -> ModuleType: return np From c9c7616df9eaaf337799e50ef9870d12d1038b89 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Nov 2023 23:04:07 +0000 Subject: [PATCH 5/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/tests/test_namedarray.py | 1 - 1 file changed, 1 deletion(-) diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index b3ae5781d99..0a7b1ceaad4 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -19,7 +19,6 @@ from numpy.typing import ArrayLike, DTypeLike, NDArray from xarray.namedarray._typing import ( - _arrayapi, _AttrsLike, _DimsLike, _DType, From 20e1d3f3e26ed0f59a63e69ba690246305c76d9a Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Mon, 4 Dec 2023 21:44:39 +0100 Subject: [PATCH 6/7] Update _typing.py --- xarray/namedarray/_typing.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index b6cd7d1fe66..af99116b08e 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -71,10 +71,11 @@ def dtype(self) -> _DType_co: _DimsLike = Union[str, Iterable[_Dim]] # https://data-apis.org/array-api/latest/API_specification/indexing.html -# TODO: np.array_api doesn't allow None for some reason, maybe they're -# recommending to use expand_dims? +# TODO: np.array_api was bugged and didn't allow (None,), but should! +# https://github.com/numpy/numpy/pull/25022 +# https://github.com/data-apis/array-api/pull/674 _IndexKey = Union[int, slice, "ellipsis"] -_IndexKeys = tuple[Union[_IndexKey], ...] +_IndexKeys = tuple[Union[_IndexKey], ...] # tuple[Union[_IndexKey, None], ...] _IndexKeyLike = Union[_IndexKey, _IndexKeys] _AttrsLike = Union[Mapping[Any, Any], None] From a370297bf77d5b7880bf76255a2f5bfdd4dab4a9 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Mon, 4 Dec 2023 22:10:32 +0100 Subject: [PATCH 7/7] Update _typing.py --- xarray/namedarray/_typing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index af99116b08e..37832daca58 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -193,7 +193,8 @@ class _arrayapi(_array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType def __getitem__( self, - key: _IndexKeyLike | Any, # TODO: Any should be _arrayapi + key: _IndexKeyLike + | Any, # TODO: Any should be _arrayapi[Any, _dtype[np.integer]] /, ) -> _arrayapi[Any, Any]: ...