diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index 4b763084e17..1fc5cacde5e 100644 --- a/xarray/namedarray/_typing.py +++ b/xarray/namedarray/_typing.py @@ -11,6 +11,7 @@ TypeVar, Union, runtime_checkable, + overload, ) import numpy as np @@ -92,10 +93,13 @@ def shape(self) -> _Shape: def dtype(self) -> _DType_co: ... - # def __array__( - # self, *, dtype: _DTypeLike[_ScalarType_co] - # ) -> NDArray[_ScalarType_co]: - # ... + @overload + def __array__(self, dtype: None = ..., /) -> np.ndarray[Any, _DType_co]: + ... + + @overload + def __array__(self, dtype: _DType, /) -> np.ndarray[Any, _DType]: + ... # Corresponds to np.typing.NDArray: diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py index 8ca631e9000..d5a7bfdac71 100644 --- a/xarray/namedarray/core.py +++ b/xarray/namedarray/core.py @@ -538,7 +538,7 @@ def _nonzero(self: T_NamedArrayInteger) -> tuple[T_NamedArrayInteger, ...]: nonzeros = np.nonzero(cast(NDArray[np.integer[Any]], self.data)) _attrs = self.attrs return tuple( - (T_NamedArrayInteger, _new(self, (dim,), nz, _attrs)) + cast(T_NamedArrayInteger, _new(self, (dim,), nz, _attrs)) for nz, dim in zip(nonzeros, self.dims) )