More typing fixes.

dcherian · Jun 13, 2024 · 8a980ef · 8a980ef
1 parent 162e34a
commit 8a980ef
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
@@ -104,7 +104,8 @@
         ReindexMethodOptions,
         Self,
         SideOptions,
-        T_Chunks,
+        T_ChunkDimFreq,
+        T_ChunksFreq,
         T_Xarray,
     )
     from xarray.core.weighted import DataArrayWeighted
@@ -1336,15 +1337,15 @@ def chunksizes(self) -> Mapping[Any, tuple[int, ...]]:
     @_deprecate_positional_args("v2023.10.0")
     def chunk(
         self,
-        chunks: T_Chunks = {},  # {} even though it's technically unsafe, is being used intentionally here (#4667)
+        chunks: T_ChunksFreq = {},  # {} 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: Any,
+        **chunks_kwargs: T_ChunkDimFreq,
     ) -> Self:
         """Coerce this array's data into a dask arrays with the given chunks.
 

diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -91,7 +91,7 @@
     QuantileMethods,
     Self,
     T_ChunkDim,
-    T_Chunks,
+    T_ChunksFreq,
     T_DataArray,
     T_DataArrayOrSet,
     T_Dataset,
@@ -159,6 +159,7 @@
         QueryParserOptions,
         ReindexMethodOptions,
         SideOptions,
+        T_ChunkDimFreq,
         T_Xarray,
     )
     from xarray.core.weighted import DatasetWeighted
@@ -2642,14 +2643,14 @@ def chunksizes(self) -> Mapping[Hashable, tuple[int, ...]]:
 
     def chunk(
         self,
-        chunks: T_Chunks = {},  # {} even though it's technically unsafe, is being used intentionally here (#4667)
+        chunks: T_ChunksFreq = {},  # {} 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: T_ChunkDim,
+        **chunks_kwargs: T_ChunkDimFreq,
     ) -> Self:
         """Coerce all arrays in this dataset into dask arrays with the given
         chunks.

diff --git a/xarray/core/types.py b/xarray/core/types.py
@@ -182,9 +182,10 @@ def copy(
 Dims = Union[str, Collection[Hashable], "ellipsis", None]
 
 # FYI in some cases we don't allow `None`, which this doesn't take account of.
-T_FreqStr: TypeAlias = str
+T_FreqStr: TypeAlias = str  # TODO: explicit list of supported pandas frequency strings
 T_ChunkDim: TypeAlias = Union[int, Literal["auto"], None, tuple[int, ...]]
 T_ChunkDimFreq: TypeAlias = Union[T_FreqStr, T_ChunkDim]
+T_ChunksFreq: TypeAlias = Union[T_ChunkDimFreq, Mapping[Any, T_ChunkDimFreq]]
 # 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]]
 T_NormalizedChunks = tuple[tuple[int, ...], ...]