Type chunkmanagers

xarray/coding/times.py

@@ -5,7 +5,7 @@
 from collections.abc import Hashable
 from datetime import datetime, timedelta
 from functools import partial
-from typing import TYPE_CHECKING, Callable, Union
+from typing import TYPE_CHECKING, Callable, Union, overload
 
 import numpy as np
 import pandas as pd
@@ -22,13 +22,17 @@
 )
 from xarray.core import indexing
 from xarray.core.common import contains_cftime_datetimes, is_np_datetime_like
-from xarray.core.duck_array_ops import asarray
+from xarray.core.duck_array_ops import asarray, ravel
 from xarray.core.formatting import first_n_items, format_timestamp, last_item
 from xarray.core.pdcompat import nanosecond_precision_timestamp
 from xarray.core.utils import emit_user_level_warning
 from xarray.core.variable import Variable
-from xarray.namedarray.parallelcompat import T_ChunkedArray, get_chunked_array_type
-from xarray.namedarray.pycompat import is_chunked_array
+from xarray.namedarray._typing import (
+    _chunkedarrayfunction_or_api,
+    chunkedduckarray,
+    duckarray,
+)
+from xarray.namedarray.parallelcompat import get_chunked_array_type
 from xarray.namedarray.utils import is_duck_dask_array
 
 try:
@@ -37,7 +41,7 @@
     cftime = None
 
 if TYPE_CHECKING:
-    from xarray.core.types import CFCalendar, T_DuckArray
+    from xarray.core.types import CFCalendar
 
     T_Name = Union[Hashable, None]
 
@@ -315,7 +319,7 @@ def decode_cf_datetime(
     cftime.num2date
     """
     num_dates = np.asarray(num_dates)
-    flat_num_dates = num_dates.ravel()
+    flat_num_dates = ravel(num_dates)
     if calendar is None:
         calendar = "standard"
 
@@ -369,7 +373,7 @@ def decode_cf_timedelta(num_timedeltas, units: str) -> np.ndarray:
     """
     num_timedeltas = np.asarray(num_timedeltas)
     units = _netcdf_to_numpy_timeunit(units)
-    result = to_timedelta_unboxed(num_timedeltas.ravel(), unit=units)
+    result = to_timedelta_unboxed(ravel(num_timedeltas), unit=units)
     return result.reshape(num_timedeltas.shape)
 
 
@@ -428,7 +432,7 @@ def infer_datetime_units(dates) -> str:
     'hours', 'minutes' or 'seconds' (the first one that can evenly divide all
     unique time deltas in `dates`)
     """
-    dates = np.asarray(dates).ravel()
+    dates = ravel(np.asarray(dates))
     if np.asarray(dates).dtype == "datetime64[ns]":
         dates = to_datetime_unboxed(dates)
         dates = dates[pd.notnull(dates)]
@@ -456,7 +460,7 @@ def infer_timedelta_units(deltas) -> str:
     {'days', 'hours', 'minutes' 'seconds'} (the first one that can evenly
     divide all unique time deltas in `deltas`)
     """
-    deltas = to_timedelta_unboxed(np.asarray(deltas).ravel())
+    deltas = to_timedelta_unboxed(ravel(np.asarray(deltas)))
     unique_timedeltas = np.unique(deltas[pd.notnull(deltas)])
     return _infer_time_units_from_diff(unique_timedeltas)
 
@@ -643,7 +647,7 @@ def encode_datetime(d):
         except TypeError:
             return np.nan if d is None else cftime.date2num(d, units, calendar)
 
-    return np.array([encode_datetime(d) for d in dates.ravel()]).reshape(dates.shape)
+    return np.array([encode_datetime(d) for d in ravel(dates)]).reshape(dates.shape)
 
 
 def cast_to_int_if_safe(num) -> np.ndarray:
@@ -700,12 +704,26 @@ def _cast_to_dtype_if_safe(num: np.ndarray, dtype: np.dtype) -> np.ndarray:
     return cast_num
 
 
+@overload
+def encode_cf_datetime(
+    dates: chunkedduckarray,
+    units: str | None = None,
+    calendar: str | None = None,
+    dtype: np.dtype | None = None,
+) -> tuple[chunkedduckarray, str, str]: ...
+@overload
+def encode_cf_datetime(
+    dates: duckarray,
+    units: str | None = None,
+    calendar: str | None = None,
+    dtype: np.dtype | None = None,
+) -> tuple[duckarray, str, str]: ...
 def encode_cf_datetime(
-    dates: T_DuckArray,  # type: ignore
+    dates: duckarray | chunkedduckarray,
     units: str | None = None,
     calendar: str | None = None,
     dtype: np.dtype | None = None,
-) -> tuple[T_DuckArray, str, str]:
+) -> tuple[duckarray | chunkedduckarray, str, str]:
     """Given an array of datetime objects, returns the tuple `(num, units,
     calendar)` suitable for a CF compliant time variable.
 
@@ -716,19 +734,19 @@ def encode_cf_datetime(
     cftime.date2num
     """
     dates = asarray(dates)
-    if is_chunked_array(dates):
+    if isinstance(dates, _chunkedarrayfunction_or_api):
         return _lazily_encode_cf_datetime(dates, units, calendar, dtype)
     else:
         return _eagerly_encode_cf_datetime(dates, units, calendar, dtype)
 
 
 def _eagerly_encode_cf_datetime(
-    dates: T_DuckArray,  # type: ignore
+    dates: duckarray,
     units: str | None = None,
     calendar: str | None = None,
     dtype: np.dtype | None = None,
     allow_units_modification: bool = True,
-) -> tuple[T_DuckArray, str, str]:
+) -> tuple[duckarray, str, str]:
     dates = asarray(dates)
 
     data_units = infer_datetime_units(dates)
@@ -753,7 +771,7 @@ def _eagerly_encode_cf_datetime(
         # Wrap the dates in a DatetimeIndex to do the subtraction to ensure
         # an OverflowError is raised if the ref_date is too far away from
         # dates to be encoded (GH 2272).
-        dates_as_index = pd.DatetimeIndex(dates.ravel())
+        dates_as_index = pd.DatetimeIndex(ravel(dates))
         time_deltas = dates_as_index - ref_date
 
         # retrieve needed units to faithfully encode to int64
@@ -806,23 +824,23 @@ def _eagerly_encode_cf_datetime(
 
 
 def _encode_cf_datetime_within_map_blocks(
-    dates: T_DuckArray,  # type: ignore
+    dates: duckarray,
     units: str,
     calendar: str,
     dtype: np.dtype,
-) -> T_DuckArray:
+) -> duckarray:
     num, *_ = _eagerly_encode_cf_datetime(
         dates, units, calendar, dtype, allow_units_modification=False
     )
     return num
 
 
 def _lazily_encode_cf_datetime(
-    dates: T_ChunkedArray,
+    dates: chunkedduckarray,
     units: str | None = None,
     calendar: str | None = None,
     dtype: np.dtype | None = None,
-) -> tuple[T_ChunkedArray, str, str]:
+) -> tuple[chunkedduckarray, str, str]:
     if calendar is None:
         # This will only trigger minor compute if dates is an object dtype array.
         calendar = infer_calendar_name(dates)
@@ -855,31 +873,43 @@ def _lazily_encode_cf_datetime(
     return num, units, calendar
 
 
+@overload
 def encode_cf_timedelta(
-    timedeltas: T_DuckArray,  # type: ignore
+    timedeltas: chunkedduckarray,
     units: str | None = None,
     dtype: np.dtype | None = None,
-) -> tuple[T_DuckArray, str]:
+) -> tuple[chunkedduckarray, str]: ...
+@overload
+def encode_cf_timedelta(
+    timedeltas: duckarray,
+    units: str | None = None,
+    dtype: np.dtype | None = None,
+) -> tuple[duckarray, str]: ...
+def encode_cf_timedelta(
+    timedeltas: chunkedduckarray | duckarray,
+    units: str | None = None,
+    dtype: np.dtype | None = None,
+) -> tuple[chunkedduckarray | duckarray, str]:
     timedeltas = asarray(timedeltas)
-    if is_chunked_array(timedeltas):
+    if isinstance(timedeltas, _chunkedarrayfunction_or_api):
         return _lazily_encode_cf_timedelta(timedeltas, units, dtype)
     else:
         return _eagerly_encode_cf_timedelta(timedeltas, units, dtype)
 
 
 def _eagerly_encode_cf_timedelta(
-    timedeltas: T_DuckArray,  # type: ignore
+    timedeltas: duckarray,
     units: str | None = None,
     dtype: np.dtype | None = None,
     allow_units_modification: bool = True,
-) -> tuple[T_DuckArray, str]:
+) -> tuple[duckarray, str]:
     data_units = infer_timedelta_units(timedeltas)
 
     if units is None:
         units = data_units
 
     time_delta = _time_units_to_timedelta64(units)
-    time_deltas = pd.TimedeltaIndex(timedeltas.ravel())
+    time_deltas = pd.TimedeltaIndex(ravel(timedeltas))
 
     # retrieve needed units to faithfully encode to int64
     needed_units = data_units
@@ -920,19 +950,21 @@ def _eagerly_encode_cf_timedelta(
 
 
 def _encode_cf_timedelta_within_map_blocks(
-    timedeltas: T_DuckArray,  # type:ignore
+    timedeltas: duckarray,
     units: str,
     dtype: np.dtype,
-) -> T_DuckArray:
+) -> duckarray:
     num, _ = _eagerly_encode_cf_timedelta(
         timedeltas, units, dtype, allow_units_modification=False
     )
     return num
 
 
 def _lazily_encode_cf_timedelta(
-    timedeltas: T_ChunkedArray, units: str | None = None, dtype: np.dtype | None = None
-) -> tuple[T_ChunkedArray, str]:
+    timedeltas: chunkedduckarray,
+    units: str | None = None,
+    dtype: np.dtype | None = None,
+) -> tuple[chunkedduckarray, str]:
     if units is None and dtype is None:
         units = "nanoseconds"
         dtype = np.dtype("int64")

xarray/coding/variables.py

@@ -179,7 +179,7 @@ def lazy_elemwise_func(array, func: Callable, dtype: np.typing.DTypeLike):
     if is_chunked_array(array):
         chunkmanager = get_chunked_array_type(array)
 
-        return chunkmanager.map_blocks(func, array, dtype=dtype)  # type: ignore[arg-type]
+        return chunkmanager.map_blocks(func, array, dtype=dtype)
     else:
         return _ElementwiseFunctionArray(array, func, dtype)
 

xarray/core/computation.py

@@ -25,6 +25,7 @@
 from xarray.core.types import Dims, T_DataArray
 from xarray.core.utils import is_dict_like, is_duck_dask_array, is_scalar, parse_dims
 from xarray.core.variable import Variable
+from xarray.namedarray._typing import chunkedduckarray
 from xarray.namedarray.parallelcompat import get_chunked_array_type
 from xarray.namedarray.pycompat import is_chunked_array
 from xarray.util.deprecation_helpers import deprecate_dims
@@ -795,6 +796,7 @@ def apply_variable_ufunc(
                     )
 
             def func(*arrays):
+                res: chunkedduckarray | tuple[chunkedduckarray, ...]
                 res = chunkmanager.apply_gufunc(
                     numpy_func,
                     signature.to_gufunc_string(exclude_dims),

xarray/core/dataset.py

@@ -163,6 +163,7 @@
         T_Xarray,
     )
     from xarray.core.weighted import DatasetWeighted
+    from xarray.namedarray._typing import duckarray
     from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint
 
 
@@ -860,7 +861,7 @@ def load(self, **kwargs) -> Self:
             chunkmanager = get_chunked_array_type(*lazy_data.values())
 
             # evaluate all the chunked arrays simultaneously
-            evaluated_data: tuple[np.ndarray[Any, Any], ...] = chunkmanager.compute(
+            evaluated_data: tuple[duckarray[Any, Any], ...] = chunkmanager.compute(
                 *lazy_data.values(), **kwargs
             )
 

xarray/core/indexing.py

@@ -1325,7 +1325,7 @@ def _arrayize_vectorized_indexer(
 
 
 def _chunked_array_with_chunks_hint(
-    array, chunks, chunkmanager: ChunkManagerEntrypoint[Any]
+    array, chunks, chunkmanager: ChunkManagerEntrypoint
 ):
     """Create a chunked array using the chunks hint for dimensions of size > 1."""
 

xarray/core/variable.py

@@ -2526,7 +2526,7 @@ def chunk(  # type: ignore[override]
         name: str | None = None,
         lock: bool | None = None,
         inline_array: bool | None = None,
-        chunked_array_type: str | ChunkManagerEntrypoint[Any] | None = None,
+        chunked_array_type: str | ChunkManagerEntrypoint | None = None,
         from_array_kwargs: Any = None,
         **chunks_kwargs: Any,
     ) -> Self:

xarray/namedarray/_typing.py

@@ -76,8 +76,13 @@ 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, ...]
+_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 +240,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 +251,11 @@ class _chunkedarrayfunction(
     @property
     def chunks(self) -> _Chunks: ...
 
+    def rechunk(
+        self,
+        chunks: _ChunksLike,
+    ) -> _chunkedarrayfunction[_ShapeType, _DType_co]: ...
+
 
 @runtime_checkable
 class _chunkedarrayapi(
@@ -260,6 +270,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)

xarray/namedarray/core.py

@@ -749,7 +749,7 @@ def sizes(self) -> dict[_Dim, _IntOrUnknown]:
     def chunk(
         self,
         chunks: int | Literal["auto"] | Mapping[Any, None | int | tuple[int, ...]] = {},
-        chunked_array_type: str | ChunkManagerEntrypoint[Any] | None = None,
+        chunked_array_type: str | ChunkManagerEntrypoint | None = None,
         from_array_kwargs: Any = None,
         **chunks_kwargs: Any,
     ) -> Self:
@@ -822,6 +822,7 @@ def chunk(
         chunkmanager = guess_chunkmanager(chunked_array_type)
 
         data_old = self._data
+        data_chunked: _chunkedarray[Any, _DType_co]
         if chunkmanager.is_chunked_array(data_old):
             data_chunked = chunkmanager.rechunk(data_old, chunks)  # type: ignore[arg-type]
         else: