Update typing imports and add compatibility for Python 3.11

xarray/namedarray/_typing.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import typing
+import sys
 from collections.abc import Hashable, Iterable, Mapping, Sequence
 from enum import Enum
 from types import ModuleType
@@ -18,15 +18,11 @@
 )
 
 import numpy as np
-from numpy.typing import NDArray
 
-if typing.TYPE_CHECKING:
-    try:
-        from dask.array.core import Array as DaskArray
-        from dask.typing import DaskCollection
-    except ImportError:
-        DaskArray = NDArray  # type: ignore
-        DaskCollection: Any = NDArray  # type: ignore
+if sys.version_info >= (3, 11):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
 
 
 # Singleton type, as per https://github.com/python/typing/pull/240
@@ -76,6 +72,10 @@ def dtype(self) -> _DType_co:
 
 _Chunks = tuple[_Shape, ...]
 _NormalizedChunks = tuple[tuple[int, ...], ...]
+# FYI in some cases we don't allow `None`, which this doesn't take account of.
+T_ChunkDim: TypeAlias = Union[int, Literal["auto"], None, tuple[int, ...]]
+# We allow the tuple form of this (though arguably we could transition to named dims only)
+T_Chunks: TypeAlias = Union[T_ChunkDim, Mapping[Any, T_ChunkDim]]
 
 _Dim = Hashable
 _Dims = tuple[_Dim, ...]

xarray/namedarray/daskmanager.py

@@ -11,8 +11,16 @@
 from xarray.namedarray.utils import is_duck_dask_array, module_available
 
 if TYPE_CHECKING:
-    from xarray.core.types import T_Chunks
-    from xarray.namedarray._typing import DaskArray, _NormalizedChunks, duckarray
+    from numpy.typing import NDArray
+
+    from xarray.namedarray._typing import T_Chunks, _NormalizedChunks, duckarray
+
+    try:
+        from dask.array.core import Array as DaskArray
+        from dask.typing import DaskCollection
+    except ImportError:
+        DaskArray = NDArray  # type: ignore
+        DaskCollection: Any = NDArray  # type: ignore
 
 
 dask_available = module_available("dask")
@@ -81,9 +89,9 @@ def array_api(self) -> Any:
     def reduction(
         self,
         arr: T_ChunkedArray,
-        func: Callable,
-        combine_func: Callable | None = None,
-        aggregate_func: Callable | None = None,
+        func: Callable[..., Any],
+        combine_func: Callable[..., Any] | None = None,
+        aggregate_func: Callable[..., Any] | None = None,
         axis: int | Sequence[int] | None = None,
         dtype: np.dtype | None = None,
         keepdims: bool = False,
@@ -102,8 +110,8 @@ def reduction(
 
     def scan(
         self,
-        func: Callable,
-        binop: Callable,
+        func: Callable[..., Any],
+        binop: Callable[..., Any],
         ident: float,
         arr: T_ChunkedArray,
         axis: int | None = None,
@@ -124,7 +132,7 @@ def scan(
 
     def apply_gufunc(
         self,
-        func: Callable,
+        func: Callable[..., Any],
         signature: str,
         *args: Any,
         axes: Sequence[tuple[int, ...]] | None = None,
@@ -156,7 +164,7 @@ def apply_gufunc(
 
     def map_blocks(
         self,
-        func: Callable,
+        func: Callable[..., Any],
         *args: Any,
         dtype: np.typing.DTypeLike | None = None,
         chunks: tuple[int, ...] | None = None,
@@ -185,14 +193,14 @@ def map_blocks(
 
     def blockwise(
         self,
-        func: Callable,
+        func: Callable[..., Any],
         out_ind: Iterable,
         *args: Any,
         # can't type this as mypy assumes args are all same type, but dask blockwise args alternate types
         name: str | None = None,
         token=None,
         dtype: np.dtype | None = None,
-        adjust_chunks: dict[Any, Callable] | None = None,
+        adjust_chunks: dict[Any, Callable[..., Any]] | None = None,
         new_axes: dict[Any, int] | None = None,
         align_arrays: bool = True,
         concatenate: bool | None = None,