From ac6fc13cf7adfbab91e1754eaa940cbcc3decdfb Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Thu, 11 Jul 2024 21:04:50 +0200 Subject: [PATCH 01/14] Update test_namedarray.py --- xarray/tests/test_namedarray.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index 3d3584448de..9c06a913921 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -79,6 +79,8 @@ def __array_namespace__(self) -> ModuleType: def check_duck_array_typevar(a: duckarray[Any, _DType]) -> duckarray[Any, _DType]: + reveal_type(a) + # Mypy checks a is valid: b: duckarray[Any, _DType] = a From 184c21176b75ef551d29234a5a815a97e1b40842 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Thu, 11 Jul 2024 22:01:18 +0200 Subject: [PATCH 02/14] Update _typing.py --- xarray/namedarray/_typing.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index b715973814f..1a169a28ff9 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -211,9 +211,7 @@ def __array_namespace__(self) -> ModuleType: ... # NamedArray can most likely use both __array_function__ and __array_namespace__: _arrayfunction_or_api = (_arrayfunction, _arrayapi) -duckarray = Union[ - _arrayfunction[_ShapeType_co, _DType_co], _arrayapi[_ShapeType_co, _DType_co] -] +duckarray = Union[_arrayfunction[_ShapeType, _DType], _arrayapi[_ShapeType, _DType]] # Corresponds to np.typing.NDArray: DuckArray = _arrayfunction[Any, np.dtype[_ScalarType_co]] From 8b649cdb53159acb0a699676a25b15b65cd75ce7 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Tue, 16 Jul 2024 22:22:44 +0200 Subject: [PATCH 03/14] Update test_namedarray.py --- xarray/tests/test_namedarray.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index 9c06a913921..3d3584448de 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -79,8 +79,6 @@ def __array_namespace__(self) -> ModuleType: def check_duck_array_typevar(a: duckarray[Any, _DType]) -> duckarray[Any, _DType]: - reveal_type(a) - # Mypy checks a is valid: b: duckarray[Any, _DType] = a From bd2b719e60841e46bad2ee3c8e4f5a91e608b784 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Thu, 18 Jul 2024 14:10:07 +0200 Subject: [PATCH 04/14] Update test_namedarray.py --- xarray/tests/test_namedarray.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index d8cea07e6ca..3d3584448de 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -180,7 +180,7 @@ def test_init(self, expected: Any) -> None: # Fail: ( ("x",), - NamedArray("time", np.array([1, 2, 3])), # type: ignore + NamedArray("time", np.array([1, 2, 3])), np.array([1, 2, 3]), True, ), @@ -341,7 +341,7 @@ def test_dims_setter( def test_duck_array_class(self) -> None: numpy_a: NDArray[np.int64] numpy_a = np.array([2.1, 4], dtype=np.dtype(np.int64)) - check_duck_array_typevar(numpy_a) # type: ignore + check_duck_array_typevar(numpy_a) masked_a: np.ma.MaskedArray[Any, np.dtype[np.int64]] masked_a = np.ma.asarray([2.1, 4], dtype=np.dtype(np.int64)) # type: ignore[no-untyped-call] @@ -560,4 +560,4 @@ def test_broadcast_to_errors( def test_warn_on_repeated_dimension_names(self) -> None: with pytest.warns(UserWarning, match="Duplicate dimension names"): - NamedArray(("x", "x"), np.arange(4).reshape(2, 2)) # type: ignore + NamedArray(("x", "x"), np.arange(4).reshape(2, 2)) From 06b8724dc94a11932a5984cbb91f6838a1d40440 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Fri, 19 Jul 2024 17:16:45 +0200 Subject: [PATCH 05/14] Update _typing.py --- xarray/namedarray/_typing.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 1a169a28ff9..13bb0f43568 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -8,6 +8,7 @@ TYPE_CHECKING, Any, Callable, + Collection, Final, Literal, Protocol, @@ -66,11 +67,11 @@ def dtype(self) -> _DType_co: ... ] # For unknown shapes Dask uses np.nan, array_api uses None: -_IntOrUnknown = int +_IntOrUnknown = Union[int, float, None] _Shape = tuple[_IntOrUnknown, ...] _ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]] -_ShapeType = TypeVar("_ShapeType", bound=Any) -_ShapeType_co = TypeVar("_ShapeType_co", bound=Any, covariant=True) +_ShapeType = TypeVar("_ShapeType", bound=_Shape) +_ShapeType_co = TypeVar("_ShapeType_co", bound=_Shape, covariant=True) _Axis = int _Axes = tuple[_Axis, ...] @@ -118,7 +119,7 @@ class _array(Protocol[_ShapeType_co, _DType_co]): """ @property - def shape(self) -> _Shape: ... + def shape(self) -> _ShapeType_co: ... @property def dtype(self) -> _DType_co: ... @@ -211,7 +212,10 @@ def __array_namespace__(self) -> ModuleType: ... # NamedArray can most likely use both __array_function__ and __array_namespace__: _arrayfunction_or_api = (_arrayfunction, _arrayapi) -duckarray = Union[_arrayfunction[_ShapeType, _DType], _arrayapi[_ShapeType, _DType]] +duckarray: TypeAlias = ( + _arrayfunction[_ShapeType, _DType] | _arrayapi[_ShapeType, _DType] +) + # Corresponds to np.typing.NDArray: DuckArray = _arrayfunction[Any, np.dtype[_ScalarType_co]] From 16b3403f450eeda0291632fdb5237762df1e979f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 19 Jul 2024 15:18:31 +0000 Subject: [PATCH 06/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/namedarray/_typing.py | 1 - 1 file changed, 1 deletion(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 13bb0f43568..42dec00f6ed 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -8,7 +8,6 @@ TYPE_CHECKING, Any, Callable, - Collection, Final, Literal, Protocol, From 0f3afddcfa7269e934be95dc6a33959f1dcb629e Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Fri, 19 Jul 2024 17:23:17 +0200 Subject: [PATCH 07/14] Update _typing.py --- xarray/namedarray/_typing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 42dec00f6ed..289bda8e036 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -66,7 +66,7 @@ def dtype(self) -> _DType_co: ... ] # For unknown shapes Dask uses np.nan, array_api uses None: -_IntOrUnknown = Union[int, float, None] +_IntOrUnknown = int _Shape = tuple[_IntOrUnknown, ...] _ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]] _ShapeType = TypeVar("_ShapeType", bound=_Shape) From ee4cd841dafafbba0ac854d501f2a4009c893115 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Fri, 19 Jul 2024 18:38:48 +0200 Subject: [PATCH 08/14] Add shapetypes --- xarray/namedarray/_typing.py | 2 +- xarray/namedarray/core.py | 12 ++++++------ xarray/tests/test_namedarray.py | 15 +++++++++------ 3 files changed, 16 insertions(+), 13 deletions(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 289bda8e036..b9a834620d7 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -212,7 +212,7 @@ def __array_namespace__(self) -> ModuleType: ... _arrayfunction_or_api = (_arrayfunction, _arrayapi) duckarray: TypeAlias = ( - _arrayfunction[_ShapeType, _DType] | _arrayapi[_ShapeType, _DType] + _arrayfunction[_ShapeType_co, _DType_co] | _arrayapi[_ShapeType_co, _DType_co] ) diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index fe47bf50533..a3db6c58a04 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -250,14 +250,14 @@ class NamedArray(NamedArrayAggregations, Generic[_ShapeType_co, _DType_co]): __slots__ = ("_data", "_dims", "_attrs") - _data: duckarray[Any, _DType_co] + _data: duckarray[_ShapeType_co, _DType_co] _dims: _Dims _attrs: dict[Any, Any] | None def __init__( self, dims: _DimsLike, - data: duckarray[Any, _DType_co], + data: duckarray[_ShapeType_co, _DType_co], attrs: _AttrsLike = None, ): self._data = data @@ -292,7 +292,7 @@ def _new( def _new( self, dims: _DimsLike | Default = _default, - data: duckarray[Any, _DType] | Default = _default, + data: duckarray[_ShapeType, _DType] | Default = _default, attrs: _AttrsLike | Default = _default, ) -> NamedArray[_ShapeType, _DType] | NamedArray[_ShapeType_co, _DType_co]: """ @@ -447,7 +447,7 @@ def dtype(self) -> _DType_co: return self._data.dtype @property - def shape(self) -> _Shape: + def shape(self) -> _ShapeType_co: """ Get the shape of the array. @@ -850,9 +850,9 @@ def to_numpy(self) -> np.ndarray[Any, Any]: # TODO an entrypoint so array libraries can choose coercion method? return to_numpy(self._data) - def as_numpy(self) -> Self: + def as_numpy(self) -> NamedArray[Any, Any]: """Coerces wrapped data into a numpy array, returning a Variable.""" - return self._replace(data=self.to_numpy()) + return self._new(data=self.to_numpy()) def reduce( self, diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index 3d3584448de..bdc55fcbbc1 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -32,20 +32,21 @@ _IntOrUnknown, _Shape, _ShapeLike, + _ShapeType, duckarray, ) class CustomArrayBase(Generic[_ShapeType_co, _DType_co]): - def __init__(self, array: duckarray[Any, _DType_co]) -> None: - self.array: duckarray[Any, _DType_co] = array + def __init__(self, array: duckarray[_ShapeType_co, _DType_co]) -> None: + self.array: duckarray[_ShapeType_co, _DType_co] = array @property def dtype(self) -> _DType_co: return self.array.dtype @property - def shape(self) -> _Shape: + def shape(self) -> _ShapeType_co: return self.array.shape @@ -78,9 +79,11 @@ def __array_namespace__(self) -> ModuleType: return np -def check_duck_array_typevar(a: duckarray[Any, _DType]) -> duckarray[Any, _DType]: +def check_duck_array_typevar( + a: duckarray[_ShapeType, _DType] +) -> duckarray[_ShapeType, _DType]: # Mypy checks a is valid: - b: duckarray[Any, _DType] = a + b: duckarray[_ShapeType, _DType] = a # Runtime check if valid: if isinstance(b, _arrayfunction_or_api): @@ -180,7 +183,7 @@ def test_init(self, expected: Any) -> None: # Fail: ( ("x",), - NamedArray("time", np.array([1, 2, 3])), + NamedArray("time", np.array([1, 2, 3], dtype=np.dtype(np.int64))), np.array([1, 2, 3]), True, ), From c7092458c384b3840fd6957ae923748ae183e3d6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 19 Jul 2024 16:39:31 +0000 Subject: [PATCH 09/14] [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 bdc55fcbbc1..c432cba0e49 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -30,7 +30,6 @@ _DType, _IndexKeyLike, _IntOrUnknown, - _Shape, _ShapeLike, _ShapeType, duckarray, From 23188b430898a192465f75068580c5de6a1bb8d0 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Fri, 19 Jul 2024 23:01:54 +0200 Subject: [PATCH 10/14] Fix __array__ missing copy parameter --- xarray/namedarray/_typing.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index b9a834620d7..b823299ddb0 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -153,13 +153,16 @@ def __getitem__( ) -> _arrayfunction[Any, _DType_co] | Any: ... @overload - def __array__(self, dtype: None = ..., /) -> np.ndarray[Any, _DType_co]: ... - + def __array__( + self, dtype: None = ..., /, *, copy: None | bool = ... + ) -> np.ndarray[Any, _DType_co]: ... @overload - def __array__(self, dtype: _DType, /) -> np.ndarray[Any, _DType]: ... + def __array__( + self, dtype: _DType, /, *, copy: None | bool = ... + ) -> np.ndarray[Any, _DType]: ... def __array__( - self, dtype: _DType | None = ..., / + self, dtype: _DType | None = ..., /, *, copy: None | bool = ... ) -> np.ndarray[Any, _DType] | np.ndarray[Any, _DType_co]: ... # TODO: Should return the same subclass but with a new dtype generic. @@ -211,9 +214,9 @@ def __array_namespace__(self) -> ModuleType: ... # NamedArray can most likely use both __array_function__ and __array_namespace__: _arrayfunction_or_api = (_arrayfunction, _arrayapi) -duckarray: TypeAlias = ( - _arrayfunction[_ShapeType_co, _DType_co] | _arrayapi[_ShapeType_co, _DType_co] -) +duckarray = Union[ + _arrayfunction[_ShapeType_co, _DType_co], _arrayapi[_ShapeType_co, _DType_co] +] # Corresponds to np.typing.NDArray: From 9471fff81ac8ae7520e4bb6828e679959aa3bc7b Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Fri, 19 Jul 2024 23:19:10 +0200 Subject: [PATCH 11/14] try reverting shapetypes --- xarray/namedarray/core.py | 12 ++++++------ xarray/tests/test_namedarray.py | 13 +++++-------- 2 files changed, 11 insertions(+), 14 deletions(-) diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index a3db6c58a04..fe47bf50533 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -250,14 +250,14 @@ class NamedArray(NamedArrayAggregations, Generic[_ShapeType_co, _DType_co]): __slots__ = ("_data", "_dims", "_attrs") - _data: duckarray[_ShapeType_co, _DType_co] + _data: duckarray[Any, _DType_co] _dims: _Dims _attrs: dict[Any, Any] | None def __init__( self, dims: _DimsLike, - data: duckarray[_ShapeType_co, _DType_co], + data: duckarray[Any, _DType_co], attrs: _AttrsLike = None, ): self._data = data @@ -292,7 +292,7 @@ def _new( def _new( self, dims: _DimsLike | Default = _default, - data: duckarray[_ShapeType, _DType] | Default = _default, + data: duckarray[Any, _DType] | Default = _default, attrs: _AttrsLike | Default = _default, ) -> NamedArray[_ShapeType, _DType] | NamedArray[_ShapeType_co, _DType_co]: """ @@ -447,7 +447,7 @@ def dtype(self) -> _DType_co: return self._data.dtype @property - def shape(self) -> _ShapeType_co: + def shape(self) -> _Shape: """ Get the shape of the array. @@ -850,9 +850,9 @@ def to_numpy(self) -> np.ndarray[Any, Any]: # TODO an entrypoint so array libraries can choose coercion method? return to_numpy(self._data) - def as_numpy(self) -> NamedArray[Any, Any]: + def as_numpy(self) -> Self: """Coerces wrapped data into a numpy array, returning a Variable.""" - return self._new(data=self.to_numpy()) + return self._replace(data=self.to_numpy()) def reduce( self, diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index c432cba0e49..2e6106c8df9 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -31,21 +31,20 @@ _IndexKeyLike, _IntOrUnknown, _ShapeLike, - _ShapeType, duckarray, ) class CustomArrayBase(Generic[_ShapeType_co, _DType_co]): - def __init__(self, array: duckarray[_ShapeType_co, _DType_co]) -> None: - self.array: duckarray[_ShapeType_co, _DType_co] = array + def __init__(self, array: duckarray[Any, _DType_co]) -> None: + self.array: duckarray[Any, _DType_co] = array @property def dtype(self) -> _DType_co: return self.array.dtype @property - def shape(self) -> _ShapeType_co: + def shape(self) -> _Shape: return self.array.shape @@ -78,11 +77,9 @@ def __array_namespace__(self) -> ModuleType: return np -def check_duck_array_typevar( - a: duckarray[_ShapeType, _DType] -) -> duckarray[_ShapeType, _DType]: +def check_duck_array_typevar(a: duckarray[Any, _DType]) -> duckarray[Any, _DType]: # Mypy checks a is valid: - b: duckarray[_ShapeType, _DType] = a + b: duckarray[Any, _DType] = a # Runtime check if valid: if isinstance(b, _arrayfunction_or_api): From 21cffba4dbfddfc1ef1f99ed26dab01b4b93e2b7 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Fri, 19 Jul 2024 23:20:41 +0200 Subject: [PATCH 12/14] try reverting shapetypes --- xarray/namedarray/_typing.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index b823299ddb0..934520dcd9c 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -69,8 +69,8 @@ def dtype(self) -> _DType_co: ... _IntOrUnknown = int _Shape = tuple[_IntOrUnknown, ...] _ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]] -_ShapeType = TypeVar("_ShapeType", bound=_Shape) -_ShapeType_co = TypeVar("_ShapeType_co", bound=_Shape, covariant=True) +_ShapeType = TypeVar("_ShapeType", bound=Any) +_ShapeType_co = TypeVar("_ShapeType_co", bound=Any, covariant=True) _Axis = int _Axes = tuple[_Axis, ...] @@ -118,7 +118,7 @@ class _array(Protocol[_ShapeType_co, _DType_co]): """ @property - def shape(self) -> _ShapeType_co: ... + def shape(self) -> _Shape: ... @property def dtype(self) -> _DType_co: ... From 03af23525a8168eec095fdfe3cc7fa3efad56a55 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Fri, 19 Jul 2024 23:21:59 +0200 Subject: [PATCH 13/14] Update test_namedarray.py --- xarray/tests/test_namedarray.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index 2e6106c8df9..7687765e659 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -30,6 +30,7 @@ _DType, _IndexKeyLike, _IntOrUnknown, + _Shape, _ShapeLike, duckarray, ) From cdde846c4371368bb677ed7f5e548660f70d1cae Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Fri, 19 Jul 2024 23:22:41 +0200 Subject: [PATCH 14/14] Update _typing.py --- xarray/namedarray/_typing.py | 1 - 1 file changed, 1 deletion(-) diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 934520dcd9c..c8d6953fc72 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -218,7 +218,6 @@ def __array_namespace__(self) -> ModuleType: ... _arrayfunction[_ShapeType_co, _DType_co], _arrayapi[_ShapeType_co, _DType_co] ] - # Corresponds to np.typing.NDArray: DuckArray = _arrayfunction[Any, np.dtype[_ScalarType_co]]