diff --git a/xarray/namedarray/_array_api.py b/xarray/namedarray/_array_api.py index 9cd064b110e..3f962b3f5a0 100644 --- a/xarray/namedarray/_array_api.py +++ b/xarray/namedarray/_array_api.py @@ -71,7 +71,7 @@ def astype( xp = x._data.__array_namespace__() return x._new(data=xp.astype(x._data, dtype, copy=copy)) - # np.astype doesn't exist yet: + # TODO: np.astype only exists in np 2: return x._new(data=x._data.astype(dtype, copy=copy)) # type: ignore[attr-defined] @@ -106,9 +106,12 @@ def imag( Size: 16B array([2., 4.]) """ - xp = _get_data_namespace(x) - out = x._new(data=xp.imag(x._data)) - return out + if isinstance(x._data, _arrayapi): + xp = x._data.__array_namespace__() + return x._new(data=xp.imag(x._data)) + + # TODO: np.imag only exists in np 2: + return x._new(data=x._data.imag) # type: ignore[attr-defined] def real( @@ -139,9 +142,12 @@ def real( Size: 16B array([1., 2.]) """ - xp = _get_data_namespace(x) - out = x._new(data=xp.real(x._data)) - return out + if isinstance(x._data, _arrayapi): + xp = x._data.__array_namespace__() + return x._new(data=xp.real(x._data)) + + # TODO: np.real only exists in np 2: + return x._new(data=x._data.real) # type: ignore[attr-defined] # %% Manipulation functions diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 95e7d7adfc3..0231fcf74e5 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -185,12 +185,6 @@ def __array_function__( kwargs: Mapping[str, Any], ) -> Any: ... - @property - def imag(self) -> _arrayfunction[_ShapeType_co, Any]: ... - - @property - def real(self) -> _arrayfunction[_ShapeType_co, Any]: ... - @runtime_checkable class _arrayapi(_array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co]): diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index cdf9eab5c8d..564dffdef91 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -565,12 +565,9 @@ def imag( -------- numpy.ndarray.imag """ - if isinstance(self._data, _arrayapi): - from xarray.namedarray._array_api import imag + from xarray.namedarray._array_api import imag - return imag(self) - - return self._new(data=self._data.imag) + return imag(self) @property def real( @@ -583,11 +580,9 @@ def real( -------- numpy.ndarray.real """ - if isinstance(self._data, _arrayapi): - from xarray.namedarray._array_api import real + from xarray.namedarray._array_api import real - return real(self) - return self._new(data=self._data.real) + return real(self) def __dask_tokenize__(self) -> object: # Use v.data, instead of v._data, in order to cope with the wrappers