Add chunks typing

xarray/namedarray/_typing.py

@@ -76,8 +76,14 @@ def dtype(self) -> _DType_co: ...
 _Axes = tuple[_Axis, ...]
 _AxisLike = Union[_Axis, _Axes]
 
-_Chunks = tuple[_Shape, ...]
-_NormalizedChunks = tuple[tuple[int, ...], ...]
+_Chunk = tuple[int, ...]
+_Chunks = tuple[_Chunk, ...]
+_NormalizedChunks = tuple[tuple[int, ...], ...]  # TODO: Same as Chunks.
+_ChunksLike = Union[
+    int, Literal["auto"], None, _Chunk, _Chunks
+]  # TODO: Literal["auto"]
+_ChunksType = TypeVar("_ChunksType", bound=_Chunks)
+
 # 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)
@@ -235,7 +241,7 @@ def chunks(self) -> _Chunks: ...
 
 @runtime_checkable
 class _chunkedarrayfunction(
-    _arrayfunction[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co]
+    _arrayfunction[_ShapeType, _DType_co], Protocol[_ShapeType, _DType_co]
 ):
     """
     Chunked duck array supporting NEP 18.
@@ -246,6 +252,11 @@ class _chunkedarrayfunction(
     @property
     def chunks(self) -> _Chunks: ...
 
+    def rechunk(
+        self,
+        chunks: _ChunksLike,
+    ) -> _chunkedarrayfunction[_ShapeType, _DType_co]: ...
+
 
 @runtime_checkable
 class _chunkedarrayapi(
@@ -260,6 +271,11 @@ class _chunkedarrayapi(
     @property
     def chunks(self) -> _Chunks: ...
 
+    def rechunk(
+        self,
+        chunks: _ChunksLike,
+    ) -> _chunkedarrayapi[_ShapeType_co, _DType_co]: ...
+
 
 # NamedArray can most likely use both __array_function__ and __array_namespace__:
 _chunkedarrayfunction_or_api = (_chunkedarrayfunction, _chunkedarrayapi)