Use strict type hinting for namedarray

pyproject.toml

@@ -93,6 +93,7 @@ module = [
   "cftime.*",
   "cubed.*",
   "cupy.*",
+  "dask.types.*",
   "fsspec.*",
   "h5netcdf.*",
   "h5py.*",
@@ -162,6 +163,41 @@ module = [
   "xarray.tests.test_weighted",
 ]
 
+# Use strict = true whenever namedarray has become standalone. In the meantime
+# don't forget to add all new files related to namedarray here:
+# ref: https://mypy.readthedocs.io/en/stable/existing_code.html#introduce-stricter-options
+[[tool.mypy.overrides]]
+# Start off with these
+warn_unused_configs = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+
+# Getting these passing should be easy
+strict_equality = true
+strict_concatenate = true
+
+# Strongly recommend enabling this one as soon as you can
+check_untyped_defs = true
+
+# These shouldn't be too much additional work, but may be tricky to
+# get passing if you use a lot of untyped libraries
+disallow_subclassing_any = true
+disallow_untyped_decorators = true
+disallow_any_generics = true
+
+# These next few are various gradations of forcing use of type annotations
+disallow_untyped_calls = true
+disallow_incomplete_defs = true
+disallow_untyped_defs = true
+
+# This one isn't too hard to get passing, but return on investment is lower
+no_implicit_reexport = true
+
+# This one can be tricky to get passing if you use a lot of untyped libraries
+warn_return_any = true
+
+module = ["xarray.namedarray.*", "xarray.tests.test_namedarray"]
+
 [tool.ruff]
 builtins = ["ellipsis"]
 exclude = [

xarray/core/parallel.py

@@ -443,7 +443,7 @@ def subset_dataset_to_block(
                 for dim in variable.dims:
                     chunk = chunk[chunk_index[dim]]
 
-                chunk_variable_task = (f"{name}-{gname}-{chunk[0]}",) + chunk_tuple
+                chunk_variable_task = (f"{name}-{gname}-{chunk[0]!r}",) + chunk_tuple
                 graph[chunk_variable_task] = (
                     tuple,
                     [variable.dims, chunk, variable.attrs],

xarray/namedarray/core.py

@@ -2,7 +2,6 @@
 
 import copy
 import math
-import sys
 import typing
 from collections.abc import Hashable, Iterable, Mapping
 
@@ -11,30 +10,38 @@
 # TODO: get rid of this after migrating this class to array API
 from xarray.core import dtypes
 from xarray.core.indexing import ExplicitlyIndexed
-from xarray.core.utils import Default, _default
 from xarray.namedarray.utils import (
+    Default,
     T_DuckArray,
+    _default,
+    is_chunked_duck_array,
     is_duck_array,
     is_duck_dask_array,
     to_0d_object_array,
 )
 
 if typing.TYPE_CHECKING:
-    T_NamedArray = typing.TypeVar("T_NamedArray", bound="NamedArray")
+    from xarray.namedarray.utils import Self  # type: ignore[attr-defined]
+
+    try:
+        from dask.typing import (
+            Graph,
+            NestedKeys,
+            PostComputeCallable,
+            PostPersistCallable,
+            SchedulerGetCallable,
+        )
+    except ImportError:
+        Graph: typing.Any  # type: ignore[no-redef]
+        NestedKeys: typing.Any  # type: ignore[no-redef]
+        SchedulerGetCallable: typing.Any  # type: ignore[no-redef]
+        PostComputeCallable: typing.Any  # type: ignore[no-redef]
+        PostPersistCallable: typing.Any  # type: ignore[no-redef]
+
+    # T_NamedArray = typing.TypeVar("T_NamedArray", bound="NamedArray")
     DimsInput = typing.Union[str, Iterable[Hashable]]
     Dims = tuple[Hashable, ...]
-
-
-try:
-    if sys.version_info >= (3, 11):
-        from typing import Self
-    else:
-        from typing_extensions import Self
-except ImportError:
-    if typing.TYPE_CHECKING:
-        raise
-    else:
-        Self: typing.Any = None
+    AttrsInput = typing.Union[Mapping[typing.Any, typing.Any], None]
 
 
 # TODO: Add tests!
@@ -46,7 +53,7 @@ def as_compatible_data(
         return typing.cast(T_DuckArray, data)
 
     if isinstance(data, np.ma.MaskedArray):
-        mask = np.ma.getmaskarray(data)
+        mask = np.ma.getmaskarray(data)  # type: ignore[no-untyped-call]
         if mask.any():
             # TODO: requires refactoring/vendoring xarray.core.dtypes and xarray.core.duck_array_ops
             raise NotImplementedError("MaskedArray is not supported yet")
@@ -74,13 +81,17 @@ class NamedArray(typing.Generic[T_DuckArray]):
     Numeric operations on this object implement array broadcasting and dimension alignment based on dimension names,
     rather than axis order."""
 
-    __slots__ = ("_dims", "_data", "_attrs")
+    __slots__ = ("_data", "_dims", "_attrs")
+
+    _data: T_DuckArray
+    _dims: Dims
+    _attrs: dict[typing.Any, typing.Any] | None
 
     def __init__(
         self,
         dims: DimsInput,
         data: T_DuckArray | np.typing.ArrayLike,
-        attrs: dict | None = None,
+        attrs: AttrsInput = None,
         fastpath: bool = False,
     ):
         """
@@ -105,9 +116,9 @@ def __init__(
 
 
         """
-        self._data: T_DuckArray = as_compatible_data(data, fastpath=fastpath)
-        self._dims: Dims = self._parse_dimensions(dims)
-        self._attrs: dict | None = dict(attrs) if attrs else None
+        self._data = as_compatible_data(data, fastpath=fastpath)
+        self._dims = self._parse_dimensions(dims)
+        self._attrs = dict(attrs) if attrs else None
 
     @property
     def ndim(self) -> int:
@@ -140,7 +151,7 @@ def __len__(self) -> int:
             raise TypeError("len() of unsized object") from exc
 
     @property
-    def dtype(self) -> np.dtype:
+    def dtype(self) -> np.dtype[typing.Any]:
         """
         Data-type of the array’s elements.
 
@@ -178,7 +189,7 @@ def nbytes(self) -> int:
         the bytes consumed based on the ``size`` and ``dtype``.
         """
         if hasattr(self._data, "nbytes"):
-            return self._data.nbytes
+            return self._data.nbytes  # type: ignore[no-any-return]
         else:
             return self.size * self.dtype.itemsize
 
@@ -208,7 +219,7 @@ def attrs(self) -> dict[typing.Any, typing.Any]:
         return self._attrs
 
     @attrs.setter
-    def attrs(self, value: Mapping) -> None:
+    def attrs(self, value: Mapping[typing.Any, typing.Any]) -> None:
         self._attrs = dict(value)
 
     def _check_shape(self, new_data: T_DuckArray) -> None:
@@ -256,43 +267,78 @@ def imag(self) -> Self:
         """
         return self._replace(data=self.data.imag)
 
-    def __dask_tokenize__(self):
-        # Use v.data, instead of v._data, in order to cope with the wrappers
-        # around NetCDF and the like
-        from dask.base import normalize_token
+    def __dask_tokenize__(self) -> Hashable | None:
+        if is_duck_dask_array(self._data):
+            # Use v.data, instead of v._data, in order to cope with the wrappers
+            # around NetCDF and the like
+            from dask.base import normalize_token
 
-        return normalize_token((type(self), self._dims, self.data, self.attrs))
+            s, d, a, attrs = type(self), self._dims, self.data, self.attrs
+            return normalize_token((s, d, a, attrs))  # type: ignore[no-any-return]
+        else:
+            return None
 
-    def __dask_graph__(self):
-        return self._data.__dask_graph__() if is_duck_dask_array(self._data) else None
+    def __dask_graph__(self) -> Graph | None:
+        if is_duck_dask_array(self._data):
+            return self._data.__dask_graph__()
+        else:
+            # TODO: Should this method just raise instead?
+            # raise NotImplementedError("Method requires self.data to be a dask array")
+            return None
 
-    def __dask_keys__(self):
-        return self._data.__dask_keys__()
+    def __dask_keys__(self) -> NestedKeys:
+        if is_duck_dask_array(self._data):
+            return self._data.__dask_keys__()
+        else:
+            raise AttributeError("Method requires self.data to be a dask array.")
 
-    def __dask_layers__(self):
-        return self._data.__dask_layers__()
+    def __dask_layers__(self) -> typing.Sequence[str]:
+        if is_duck_dask_array(self._data):
+            return self._data.__dask_layers__()
+        else:
+            raise AttributeError("Method requires self.data to be a dask array.")
 
     @property
-    def __dask_optimize__(self) -> typing.Callable:
-        return self._data.__dask_optimize__
+    def __dask_optimize__(
+        self,
+    ) -> typing.Callable[..., dict[typing.Any, typing.Any]]:
+        if is_duck_dask_array(self._data):
+            return self._data.__dask_optimize__()  # type: ignore[no-any-return]
+        else:
+            raise AttributeError("Method requires self.data to be a dask array.")
 
     @property
-    def __dask_scheduler__(self) -> typing.Callable:
-        return self._data.__dask_scheduler__
+    def __dask_scheduler__(self) -> staticmethod[SchedulerGetCallable]:
+        if is_duck_dask_array(self._data):
+            return self._data.__dask_scheduler__()  # type: ignore[no-any-return]
+        else:
+            raise AttributeError("Method requires self.data to be a dask array.")
 
     def __dask_postcompute__(
         self,
-    ) -> tuple[typing.Callable, tuple[typing.Any, ...]]:
-        array_func, array_args = self._data.__dask_postcompute__()
-        return self._dask_finalize, (array_func,) + array_args
+    ) -> tuple[PostComputeCallable, tuple[typing.Any, ...]]:
+        if is_duck_dask_array(self._data):
+            array_func, array_args = self._data.__dask_postcompute__()  # type: ignore[no-untyped-call]
+            return self._dask_finalize, (array_func,) + array_args
+        else:
+            raise AttributeError("Method requires self.data to be a dask array.")
 
     def __dask_postpersist__(
         self,
-    ) -> tuple[typing.Callable, tuple[typing.Any, ...]]:
-        array_func, array_args = self._data.__dask_postpersist__()
-        return self._dask_finalize, (array_func,) + array_args
+    ) -> tuple[PostPersistCallable, tuple[typing.Any, ...]]:
+        if is_duck_dask_array(self._data):
+            array_func, array_args = self._data.__dask_postpersist__()  # type: ignore[no-untyped-call]
+            return self._dask_finalize, (array_func,) + array_args
+        else:
+            raise AttributeError("Method requires self.data to be a dask array.")
 
-    def _dask_finalize(self, results, array_func, *args, **kwargs) -> Self:
+    def _dask_finalize(
+        self,
+        results: T_DuckArray,
+        array_func: typing.Callable[..., T_DuckArray],
+        *args: typing.Any,
+        **kwargs: typing.Any,
+    ) -> Self:
         data = array_func(results, *args, **kwargs)
         return type(self)(self._dims, data, attrs=self._attrs)
 
@@ -308,7 +354,13 @@ def chunks(self) -> tuple[tuple[int, ...], ...] | None:
         NamedArray.chunksizes
         xarray.unify_chunks
         """
-        return getattr(self._data, "chunks", None)
+        data = self._data
+        # reveal_type(data)
+        if is_chunked_duck_array(data):
+            # reveal_type(data)
+            return data.chunks
+        else:
+            return None
 
     @property
     def chunksizes(
@@ -328,8 +380,9 @@ def chunksizes(
         NamedArray.chunks
         xarray.unify_chunks
         """
-        if hasattr(self._data, "chunks"):
-            return dict(zip(self.dims, self.data.chunks))
+        data = self._data
+        if is_chunked_duck_array(data):
+            return dict(zip(self.dims, data.chunks))
         else:
             return {}
 
@@ -338,7 +391,12 @@ def sizes(self) -> dict[Hashable, int]:
         """Ordered mapping from dimension names to lengths."""
         return dict(zip(self.dims, self.shape))
 
-    def _replace(self, dims=_default, data=_default, attrs=_default) -> Self:
+    def _replace(
+        self,
+        dims: DimsInput | Default = _default,
+        data: T_DuckArray | np.typing.ArrayLike | Default = _default,
+        attrs: AttrsInput | Default = _default,
+    ) -> Self:
         if dims is _default:
             dims = copy.copy(self._dims)
         if data is _default:
@@ -415,7 +473,7 @@ def _nonzero(self) -> tuple[Self, ...]:
     def _as_sparse(
         self,
         sparse_format: str | Default = _default,
-        fill_value=dtypes.NA,
+        fill_value: typing.Any = dtypes.NA,
     ) -> Self:
         """
         use sparse-array as backend.