more typing

Author: andersy005

xarray/backends/common.py

@@ -13,8 +13,8 @@
 from xarray.conventions import cf_encoder
 from xarray.core import indexing
 from xarray.core.utils import FrozenDict, NdimSizeLenMixin, is_remote_uri
+from xarray.namedarray._typing import _chunkedarrayfunction_or_api
 from xarray.namedarray.parallelcompat import get_chunked_array_type
-from xarray.namedarray.pycompat import is_chunked_array
 
 if TYPE_CHECKING:
     from io import BufferedIOBase
@@ -232,7 +232,7 @@ def __init__(self, lock=None):
         self.lock = lock
 
     def add(self, source, target, region=None):
-        if is_chunked_array(source):
+        if isinstance(source, _chunkedarrayfunction_or_api):
             self.sources.append(source)
             self.targets.append(target)
             self.regions.append(region)

xarray/namedarray/parallelcompat.py

@@ -19,7 +19,10 @@
 if TYPE_CHECKING:
     from xarray.namedarray._typing import (
         _Chunks,
+        _DType,
+        _DTypeLike,
         _NormalizedChunks,
+        _ShapeType,
         duckarray,
     )
 
@@ -220,9 +223,9 @@ def chunks(self, data: T_ChunkedArray) -> _NormalizedChunks:
     def normalize_chunks(
         self,
         chunks: _Chunks | _NormalizedChunks,
-        shape: tuple[int, ...] | None = None,
+        shape: _ShapeType | None = None,
         limit: int | None = None,
-        dtype: np.dtype | None = None,
+        dtype: _DType | None = None,
         previous_chunks: _NormalizedChunks | None = None,
     ) -> _NormalizedChunks:
         """
@@ -359,7 +362,7 @@ def reduction(
         combine_func: Callable | None = None,
         aggregate_func: Callable | None = None,
         axis: int | Sequence[int] | None = None,
-        dtype: np.dtype | None = None,
+        dtype: _DTypeLike | None = None,
         keepdims: bool = False,
     ) -> T_ChunkedArray:
         """
@@ -409,7 +412,7 @@ def apply_gufunc(
         *args: Any,
         axes: Sequence[tuple[int, ...]] | None = None,
         keepdims: bool = False,
-        output_dtypes: Sequence[np.typing.DTypeLike] | None = None,
+        output_dtypes: Sequence[_DTypeLike] | None = None,
         vectorize: bool | None = None,
         **kwargs,
     ):
@@ -492,7 +495,7 @@ def map_blocks(
         self,
         func: Callable,
         *args: Any,
-        dtype: np.typing.DTypeLike | None = None,
+        dtype: _DTypeLike | None = None,
         chunks: tuple[int, ...] | None = None,
         drop_axis: int | Sequence[int] | None = None,
         new_axis: int | Sequence[int] | None = None,

xarray/namedarray/pycompat.py

@@ -86,17 +86,17 @@ def mod_version(mod: ModType) -> Version:
     return _get_cached_duck_array_module(mod).version
 
 
-def is_chunked_array(x) -> bool:
+def is_chunked_array(x: duckarray[Any, Any]) -> bool:
     return is_duck_dask_array(x) or (
         isinstance(x, _arrayfunction_or_api) and hasattr(x, "chunks")
     )
 
 
-def is_0d_dask_array(x) -> bool:
+def is_0d_dask_array(x: duckarray[Any, Any]) -> bool:
     return is_duck_dask_array(x) and is_scalar(x)
 
 
-def to_numpy(data) -> np.ndarray[_ShapeType, _DType]:
+def to_numpy(data: duckarray[Any, Any]) -> np.ndarray[_ShapeType, _DType]:
     from xarray.core.indexing import ExplicitlyIndexed
     from xarray.namedarray.parallelcompat import get_chunked_array_type