fixes

pydata · Oct 14, 2023 · 1f52099 · 1f52099
1 parent aecdff4
commit 1f52099
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git 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
@@ -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)
         )