diff --git a/xarray/namedarray/_typing.py b/xarray/namedarray/_typing.py index b715973814f..66519ffd341 100644 --- a/xarray/namedarray/_typing.py +++ b/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)