From ea06c6f2983a31186fe3bb7495f3505dbb56d630 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Wed, 25 Sep 2024 15:24:43 +0200 Subject: [PATCH] Update __array__ signatures with copy (#9529) * Update __array__ with copy * Update common.py * Update indexing.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * copy only available from np2 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Raise if copy=false * Update groupby.py * Update test_namedarray.py * Update pyproject.toml --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- pyproject.toml | 1 - xarray/core/common.py | 2 +- xarray/core/datatree.py | 5 +++- xarray/core/groupby.py | 6 ++++- xarray/core/indexing.py | 43 +++++++++++++++++++++------------ xarray/namedarray/_typing.py | 6 ++--- xarray/tests/arrays.py | 12 ++++++--- xarray/tests/test_assertions.py | 6 +++-- xarray/tests/test_formatting.py | 4 ++- xarray/tests/test_namedarray.py | 11 +++++++-- 10 files changed, 65 insertions(+), 31 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 35522d82edf..0078a346b75 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -323,7 +323,6 @@ filterwarnings = [ "default:Using a non-tuple sequence for multidimensional indexing is deprecated:FutureWarning", "default:Duplicate dimension names present:UserWarning:xarray.namedarray.core", "default:::xarray.tests.test_strategies", # TODO: remove once we know how to deal with a changed signature in protocols - "ignore:__array__ implementation doesn't accept a copy keyword, so passing copy=False failed.", ] log_cli_level = "INFO" diff --git a/xarray/core/common.py b/xarray/core/common.py index e4c61a1bc12..9a6807faad2 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -163,7 +163,7 @@ def __complex__(self: Any) -> complex: return complex(self.values) def __array__( - self: Any, dtype: DTypeLike | None = None, copy: bool | None = None + self: Any, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None ) -> np.ndarray: if not copy: if np.lib.NumpyVersion(np.__version__) >= "2.0.0": diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index bd583ac86cb..65436be9038 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -55,6 +55,7 @@ from xarray.core.dataset import calculate_dimensions if TYPE_CHECKING: + import numpy as np import pandas as pd from xarray.core.datatree_io import T_DataTreeNetcdfEngine, T_DataTreeNetcdfTypes @@ -737,7 +738,9 @@ def __bool__(self) -> bool: def __iter__(self) -> Iterator[str]: return itertools.chain(self._data_variables, self._children) # type: ignore[arg-type] - def __array__(self, dtype=None, copy=None): + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray: raise TypeError( "cannot directly convert a DataTree into a " "numpy array. Instead, create an xarray.DataArray " diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index 58971435018..92f0572d37a 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -193,7 +193,11 @@ def values(self) -> range: def data(self) -> range: return range(self.size) - def __array__(self) -> np.ndarray: + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray: + if copy is False: + raise NotImplementedError(f"An array copy is necessary, got {copy = }.") return np.arange(self.size) @property diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 67912908a2b..08b1d0be290 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -13,6 +13,7 @@ import numpy as np import pandas as pd +from packaging.version import Version from xarray.core import duck_array_ops from xarray.core.nputils import NumpyVIndexAdapter @@ -505,9 +506,14 @@ class ExplicitlyIndexed: __slots__ = () - def __array__(self, dtype: np.typing.DTypeLike = None) -> np.ndarray: + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray: # Leave casting to an array up to the underlying array type. - return np.asarray(self.get_duck_array(), dtype=dtype) + if Version(np.__version__) >= Version("2.0.0"): + return np.asarray(self.get_duck_array(), dtype=dtype, copy=copy) + else: + return np.asarray(self.get_duck_array(), dtype=dtype) def get_duck_array(self): return self.array @@ -520,11 +526,6 @@ def get_duck_array(self): key = BasicIndexer((slice(None),) * self.ndim) return self[key] - def __array__(self, dtype: np.typing.DTypeLike = None) -> np.ndarray: - # This is necessary because we apply the indexing key in self.get_duck_array() - # Note this is the base class for all lazy indexing classes - return np.asarray(self.get_duck_array(), dtype=dtype) - def _oindex_get(self, indexer: OuterIndexer): raise NotImplementedError( f"{self.__class__.__name__}._oindex_get method should be overridden" @@ -570,8 +571,13 @@ def __init__(self, array, indexer_cls: type[ExplicitIndexer] = BasicIndexer): self.array = as_indexable(array) self.indexer_cls = indexer_cls - def __array__(self, dtype: np.typing.DTypeLike = None) -> np.ndarray: - return np.asarray(self.get_duck_array(), dtype=dtype) + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray: + if Version(np.__version__) >= Version("2.0.0"): + return np.asarray(self.get_duck_array(), dtype=dtype, copy=copy) + else: + return np.asarray(self.get_duck_array(), dtype=dtype) def get_duck_array(self): return self.array.get_duck_array() @@ -830,9 +836,6 @@ def __init__(self, array): def _ensure_cached(self): self.array = as_indexable(self.array.get_duck_array()) - def __array__(self, dtype: np.typing.DTypeLike = None) -> np.ndarray: - return np.asarray(self.get_duck_array(), dtype=dtype) - def get_duck_array(self): self._ensure_cached() return self.array.get_duck_array() @@ -1674,7 +1677,9 @@ def __init__(self, array: pd.Index, dtype: DTypeLike = None): def dtype(self) -> np.dtype: return self._dtype - def __array__(self, dtype: DTypeLike = None) -> np.ndarray: + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray: if dtype is None: dtype = self.dtype array = self.array @@ -1682,7 +1687,11 @@ def __array__(self, dtype: DTypeLike = None) -> np.ndarray: with suppress(AttributeError): # this might not be public API array = array.astype("object") - return np.asarray(array.values, dtype=dtype) + + if Version(np.__version__) >= Version("2.0.0"): + return np.asarray(array.values, dtype=dtype, copy=copy) + else: + return np.asarray(array.values, dtype=dtype) def get_duck_array(self) -> np.ndarray: return np.asarray(self) @@ -1831,7 +1840,9 @@ def __init__( super().__init__(array, dtype) self.level = level - def __array__(self, dtype: DTypeLike = None) -> np.ndarray: + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray: if dtype is None: dtype = self.dtype if self.level is not None: @@ -1839,7 +1850,7 @@ def __array__(self, dtype: DTypeLike = None) -> np.ndarray: self.array.get_level_values(self.level).values, dtype=dtype ) else: - return super().__array__(dtype) + return super().__array__(dtype, copy=copy) def _convert_scalar(self, item): if isinstance(item, tuple) and self.level is not None: diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index a7d7ed7994f..90c442d2e1f 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -153,15 +153,15 @@ def __getitem__( @overload def __array__( - self, dtype: None = ..., /, *, copy: None | bool = ... + self, dtype: None = ..., /, *, copy: bool | None = ... ) -> np.ndarray[Any, _DType_co]: ... @overload def __array__( - self, dtype: _DType, /, *, copy: None | bool = ... + self, dtype: _DType, /, *, copy: bool | None = ... ) -> np.ndarray[Any, _DType]: ... def __array__( - self, dtype: _DType | None = ..., /, *, copy: None | bool = ... + self, dtype: _DType | None = ..., /, *, copy: bool | None = ... ) -> np.ndarray[Any, _DType] | np.ndarray[Any, _DType_co]: ... # TODO: Should return the same subclass but with a new dtype generic. diff --git a/xarray/tests/arrays.py b/xarray/tests/arrays.py index 4e1e31cfa49..7373b6c75ab 100644 --- a/xarray/tests/arrays.py +++ b/xarray/tests/arrays.py @@ -24,7 +24,9 @@ def __init__(self, array): def get_duck_array(self): raise UnexpectedDataAccess("Tried accessing data") - def __array__(self, dtype: np.typing.DTypeLike = None): + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray: raise UnexpectedDataAccess("Tried accessing data") def __getitem__(self, key): @@ -49,7 +51,9 @@ def __init__(self, array: np.ndarray): def __getitem__(self, key): return type(self)(self.array[key]) - def __array__(self, dtype: np.typing.DTypeLike = None): + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray: raise UnexpectedDataAccess("Tried accessing data") def __array_namespace__(self): @@ -140,7 +144,9 @@ def __repr__(self: Any) -> str: def get_duck_array(self): raise UnexpectedDataAccess("Tried accessing data") - def __array__(self, dtype: np.typing.DTypeLike = None): + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray: raise UnexpectedDataAccess("Tried accessing data") def __getitem__(self, key) -> "ConcatenatableArray": diff --git a/xarray/tests/test_assertions.py b/xarray/tests/test_assertions.py index 3e1ce0ea266..ec4b39aaab6 100644 --- a/xarray/tests/test_assertions.py +++ b/xarray/tests/test_assertions.py @@ -173,9 +173,11 @@ def dims(self): warnings.warn("warning in test", stacklevel=2) return super().dims - def __array__(self, dtype=None, copy=None): + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray: warnings.warn("warning in test", stacklevel=2) - return super().__array__() + return super().__array__(dtype, copy=copy) a = WarningVariable("x", [1]) b = WarningVariable("x", [2]) diff --git a/xarray/tests/test_formatting.py b/xarray/tests/test_formatting.py index 4123b3e8aee..688f41a7f92 100644 --- a/xarray/tests/test_formatting.py +++ b/xarray/tests/test_formatting.py @@ -942,7 +942,9 @@ def test_lazy_array_wont_compute() -> None: from xarray.core.indexing import LazilyIndexedArray class LazilyIndexedArrayNotComputable(LazilyIndexedArray): - def __array__(self, dtype=None, copy=None): + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray: raise NotImplementedError("Computing this array is not possible.") arr = LazilyIndexedArrayNotComputable(np.array([1, 2])) diff --git a/xarray/tests/test_namedarray.py b/xarray/tests/test_namedarray.py index 8ccf8c541b7..5e4a39ada73 100644 --- a/xarray/tests/test_namedarray.py +++ b/xarray/tests/test_namedarray.py @@ -8,6 +8,7 @@ import numpy as np import pytest +from packaging.version import Version from xarray.core.indexing import ExplicitlyIndexed from xarray.namedarray._typing import ( @@ -53,8 +54,14 @@ def shape(self) -> _Shape: class CustomArray( CustomArrayBase[_ShapeType_co, _DType_co], Generic[_ShapeType_co, _DType_co] ): - def __array__(self) -> np.ndarray[Any, np.dtype[np.generic]]: - return np.array(self.array) + def __array__( + self, dtype: np.typing.DTypeLike = None, /, *, copy: bool | None = None + ) -> np.ndarray[Any, np.dtype[np.generic]]: + + if Version(np.__version__) >= Version("2.0.0"): + return np.asarray(self.array, dtype=dtype, copy=copy) + else: + return np.asarray(self.array, dtype=dtype) class CustomArrayIndexable(