Fix & normalize typing for chunks (#8247)

* Fix & normalize typing for chunks

I noticed that `"auto"` wasn't allowed as a value in a dict. So this normalizes all chunk types, and defines the mapping as containing the inner type.

Allows removing some ignores (though also adds one).

One question — not necessary to answer now — is whether we should allow a tuple of definitions, for each dimension. Generally we use names, which helps prevent mistakes, and allows us to be less concerned about dimension ordering.

Author: max-sixty

xarray/core/dataarray.py

@@ -111,6 +111,7 @@
         ReindexMethodOptions,
         Self,
         SideOptions,
+        T_Chunks,
         T_Xarray,
     )
     from xarray.core.weighted import DataArrayWeighted
@@ -1288,13 +1289,7 @@ def chunksizes(self) -> Mapping[Any, tuple[int, ...]]:
 
     def chunk(
         self,
-        chunks: (
-            int
-            | Literal["auto"]
-            | tuple[int, ...]
-            | tuple[tuple[int, ...], ...]
-            | Mapping[Any, None | int | tuple[int, ...]]
-        ) = {},  # {} even though it's technically unsafe, is being used intentionally here (#4667)
+        chunks: T_Chunks = {},  # {} even though it's technically unsafe, is being used intentionally here (#4667)
         name_prefix: str = "xarray-",
         token: str | None = None,
         lock: bool = False,
@@ -1362,7 +1357,7 @@ def chunk(
 
         if isinstance(chunks, (float, str, int)):
             # ignoring type; unclear why it won't accept a Literal into the value.
-            chunks = dict.fromkeys(self.dims, chunks)  # type: ignore
+            chunks = dict.fromkeys(self.dims, chunks)
         elif isinstance(chunks, (tuple, list)):
             chunks = dict(zip(self.dims, chunks))
         else:

xarray/core/dataset.py

@@ -93,7 +93,14 @@
     is_duck_array,
     is_duck_dask_array,
 )
-from xarray.core.types import QuantileMethods, Self, T_DataArrayOrSet, T_Dataset
+from xarray.core.types import (
+    QuantileMethods,
+    Self,
+    T_ChunkDim,
+    T_Chunks,
+    T_DataArrayOrSet,
+    T_Dataset,
+)
 from xarray.core.utils import (
     Default,
     Frozen,
@@ -1478,7 +1485,7 @@ def __iter__(self) -> Iterator[Hashable]:
     if TYPE_CHECKING:
         # needed because __getattr__ is returning Any and otherwise
         # this class counts as part of the SupportsArray Protocol
-        __array__ = None
+        __array__ = None  # type: ignore[var-annotated,unused-ignore]
 
     else:
 
@@ -2569,16 +2576,14 @@ def chunksizes(self) -> Mapping[Hashable, tuple[int, ...]]:
 
     def chunk(
         self,
-        chunks: (
-            int | Literal["auto"] | Mapping[Any, None | int | str | tuple[int, ...]]
-        ) = {},  # {} even though it's technically unsafe, is being used intentionally here (#4667)
+        chunks: T_Chunks = {},  # {} even though it's technically unsafe, is being used intentionally here (#4667)
         name_prefix: str = "xarray-",
         token: str | None = None,
         lock: bool = False,
         inline_array: bool = False,
         chunked_array_type: str | ChunkManagerEntrypoint | None = None,
         from_array_kwargs=None,
-        **chunks_kwargs: None | int | str | tuple[int, ...],
+        **chunks_kwargs: T_ChunkDim,
     ) -> Self:
         """Coerce all arrays in this dataset into dask arrays with the given
         chunks.
@@ -2637,8 +2642,9 @@ def chunk(
             )
             chunks = {}
 
-        if isinstance(chunks, (Number, str, int)):
-            chunks = dict.fromkeys(self.dims, chunks)
+        if not isinstance(chunks, Mapping):
+            # We need to ignore since mypy doesn't recognize this can't be `None`
+            chunks = dict.fromkeys(self.dims, chunks)  # type: ignore[arg-type]
         else:
             chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk")
 

xarray/core/types.py

@@ -19,9 +19,9 @@
 
 try:
     if sys.version_info >= (3, 11):
-        from typing import Self
+        from typing import Self, TypeAlias
     else:
-        from typing_extensions import Self
+        from typing_extensions import Self, TypeAlias
 except ImportError:
     if TYPE_CHECKING:
         raise
@@ -183,7 +183,12 @@ def copy(
 Dims = Union[str, Iterable[Hashable], "ellipsis", None]
 OrderedDims = Union[str, Sequence[Union[Hashable, "ellipsis"]], "ellipsis", None]
 
-T_Chunks = Union[int, dict[Any, Any], Literal["auto"], None]
+# 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], tuple[T_ChunkDim, ...]
+]
 T_NormalizedChunks = tuple[tuple[int, ...], ...]
 
 DataVars = Mapping[Any, Any]

xarray/core/variable.py

@@ -1035,7 +1035,7 @@ def chunk(
 
         data_old = self._data
         if chunkmanager.is_chunked_array(data_old):
-            data_chunked = chunkmanager.rechunk(data_old, chunks)  # type: ignore[arg-type]
+            data_chunked = chunkmanager.rechunk(data_old, chunks)
         else:
             if isinstance(data_old, indexing.ExplicitlyIndexed):
                 # Unambiguously handle array storage backends (like NetCDF4 and h5py)
@@ -1057,7 +1057,7 @@ def chunk(
 
             data_chunked = chunkmanager.from_array(
                 ndata,
-                chunks,  # type: ignore[arg-type]
+                chunks,
                 **_from_array_kwargs,
             )