more typing updates

Author: andersy005

xarray/namedarray/daskmanager.py

@@ -12,6 +12,7 @@
 
 if TYPE_CHECKING:
     from xarray.core.types import DaskArray, T_Chunks, T_NormalizedChunks
+    from xarray.namedarray._typing import duckarray
 
 
 dask_available = module_available("dask")
@@ -28,7 +29,7 @@ def __init__(self) -> None:
 
         self.array_cls = Array
 
-    def is_chunked_array(self, data: Any) -> bool:
+    def is_chunked_array(self, data: duckarray[Any, Any]) -> bool:
         return is_duck_dask_array(data)
 
     def chunks(self, data: DaskArray) -> T_NormalizedChunks:

xarray/namedarray/parallelcompat.py

@@ -17,7 +17,11 @@
 from xarray.namedarray.pycompat import is_chunked_array
 
 if TYPE_CHECKING:
-    from xarray.namedarray._typing import T_DuckArray, _Chunks, _NormalizedChunks
+    from xarray.namedarray._typing import (
+        _Chunks,
+        _NormalizedChunks,
+        duckarray,
+    )
 
 T_ChunkedArray = TypeVar("T_ChunkedArray")
 
@@ -168,7 +172,7 @@ def __init__(self) -> None:
         """Used to set the array_cls attribute at import time."""
         raise NotImplementedError()
 
-    def is_chunked_array(self, data: Any) -> bool:
+    def is_chunked_array(self, data: duckarray[Any, Any]) -> bool:
         """
         Check if the given object is an instance of this type of chunked array.
 
@@ -251,7 +255,7 @@ def normalize_chunks(
 
     @abstractmethod
     def from_array(
-        self, data: T_DuckArray | np.typing.ArrayLike, chunks: _Chunks, **kwargs
+        self, data: duckarray[Any, Any], chunks: _Chunks, **kwargs
     ) -> T_ChunkedArray:
         """
         Create a chunked array from a non-chunked numpy-like array.

xarray/namedarray/pycompat.py

@@ -16,6 +16,7 @@
 if TYPE_CHECKING:
     ModType = Literal["dask", "pint", "cupy", "sparse", "cubed"]
     DuckArrayTypes = tuple[type[Any], ...]  # TODO: improve this? maybe Generic
+    from xarray.namedarray._typing import _DType, _ShapeType, duckarray
 
 
 class DuckArrayModule:
@@ -91,11 +92,11 @@ def is_chunked_array(x) -> bool:
     )
 
 
-def is_0d_dask_array(x):
+def is_0d_dask_array(x) -> bool:
     return is_duck_dask_array(x) and is_scalar(x)
 
 
-def to_numpy(data) -> np.ndarray:
+def to_numpy(data) -> np.ndarray[_ShapeType, _DType]:
     from xarray.core.indexing import ExplicitlyIndexed
     from xarray.namedarray.parallelcompat import get_chunked_array_type
 
@@ -118,7 +119,7 @@ def to_numpy(data) -> np.ndarray:
     return data
 
 
-def to_duck_array(data):
+def to_duck_array(data) -> duckarray[_ShapeType, _DType]:
     from xarray.core.indexing import ExplicitlyIndexed
 
     if isinstance(data, ExplicitlyIndexed):

xarray/namedarray/utils.py

@@ -129,11 +129,11 @@ def either_dict_or_kwargs(
 
 
 def consolidate_dask_from_array_kwargs(
-    from_array_kwargs: dict,
+    from_array_kwargs: dict[Any, Any],
     name: str | None = None,
     lock: bool | None = None,
     inline_array: bool | None = None,
-) -> dict:
+) -> dict[Any, Any]:
     """
     Merge dask-specific kwargs with arbitrary from_array_kwargs dict.
 
@@ -166,12 +166,12 @@ def consolidate_dask_from_array_kwargs(
 
 
 def _resolve_doubly_passed_kwarg(
-    kwargs_dict: dict,
+    kwargs_dict: dict[Any, Any],
     kwarg_name: str,
     passed_kwarg_value: str | bool | None,
     default: bool | None,
     err_msg_dict_name: str,
-) -> dict:
+) -> dict[Any, Any]:
     # if in kwargs_dict but not passed explicitly then just pass kwargs_dict through unaltered
     if kwarg_name in kwargs_dict and passed_kwarg_value is None:
         pass