#142 Merge pull request from astropenguin/astropenguin/issue141

Update protocol for public type hints

Author: astropenguin

tests/test_typing.py

@@ -9,8 +9,8 @@
 
 # submodules
 from xarray_dataclasses.typing import (
-    ArrayLike,
     Attr,
+    Collection,
     Coord,
     Data,
     Name,
@@ -34,21 +34,21 @@
     (Tuple[()], ()),
     (Tuple[X], ("x",)),
     (Tuple[X, Y], ("x", "y")),
-    (ArrayLike[X, Any], ("x",)),
-    (ArrayLike[Tuple[()], Any], ()),
-    (ArrayLike[Tuple[X], Any], ("x",)),
-    (ArrayLike[Tuple[X, Y], Any], ("x", "y")),
+    (Collection[X, Any], ("x",)),
+    (Collection[Tuple[()], Any], ()),
+    (Collection[Tuple[X], Any], ("x",)),
+    (Collection[Tuple[X, Y], Any], ("x", "y")),
 ]
 
 testdata_dtype = [
     (Any, None),
     (NoneType, None),
     (Int64, "int64"),
     (int, "int"),
-    (ArrayLike[Any, Any], None),
-    (ArrayLike[Any, NoneType], None),
-    (ArrayLike[Any, Int64], "int64"),
-    (ArrayLike[Any, int], "int"),
+    (Collection[Any, Any], None),
+    (Collection[Any, NoneType], None),
+    (Collection[Any, Int64], "int64"),
+    (Collection[Any, int], "int"),
 ]
 
 testdata_field_type = [

xarray_dataclasses/datamodel.py

@@ -14,7 +14,6 @@
 
 # submodules
 from .typing import (
-    ArrayLike,
     DataClass,
     DataType,
     Dims,
@@ -254,27 +253,27 @@ def get_typedarray(
         DataArray object with given dims and dtype.
 
     """
-    if isinstance(data, ArrayLike):
-        array = cast(np.ndarray, data)
-    else:
-        array = np.asarray(data)
+    try:
+        data.__array__
+    except AttributeError:
+        data = np.asarray(data)
 
     if dtype is not None:
-        array = array.astype(dtype, copy=False)
+        data = data.astype(dtype, copy=False)
 
-    if array.ndim == len(dims):
-        dataarray = xr.DataArray(array, dims=dims)
-    elif array.ndim == 0 and reference is not None:
-        dataarray = xr.DataArray(array)
+    if data.ndim == len(dims):
+        dataarray = xr.DataArray(data, dims=dims)
+    elif data.ndim == 0 and reference is not None:
+        dataarray = xr.DataArray(data)
     else:
         raise ValueError(
             "Could not create a DataArray object from data. "
-            f"Mismatch between shape {array.shape} and dims {dims}."
+            f"Mismatch between shape {data.shape} and dims {dims}."
         )
 
     if reference is None:
         return dataarray
-
-    diff_dims = set(reference.dims) - set(dims)
-    subspace = reference.isel({dim: 0 for dim in diff_dims})
-    return dataarray.broadcast_like(subspace)
+    else:
+        ddims = set(reference.dims) - set(dims)
+        reference = reference.isel({dim: 0 for dim in ddims})
+        return dataarray.broadcast_like(reference)

xarray_dataclasses/typing.py

@@ -23,6 +23,7 @@
 from typing import (
     Any,
     ClassVar,
+    Collection,
     Dict,
     Hashable,
     Optional,
@@ -44,7 +45,6 @@
     get_args,
     get_origin,
     get_type_hints,
-    runtime_checkable,
 )
 
 
@@ -65,23 +65,16 @@
 Sizes = Dict[str, int]
 
 
-@runtime_checkable
-class ArrayLike(Protocol[TDims, TDtype]):
-    """Type hint for array-like objects."""
+class Labeled(Protocol[TDims]):
+    """Type hint for labeled objects."""
 
-    def astype(self: T, dtype: Any) -> T:
-        """Method to convert data type of the object."""
-        ...
+    pass
 
-    @property
-    def ndim(self) -> int:
-        """Number of dimensions of the object."""
-        ...
 
-    @property
-    def shape(self) -> Tuple[int, ...]:
-        """Shape of the object."""
-        ...
+class Collection(Labeled[TDims], Collection[TDtype], Protocol):
+    """Type hint for labeled collection objects."""
+
+    pass
 
 
 class DataClass(Protocol[P]):
@@ -138,7 +131,7 @@ class Image(AsDataArray):
 
 """
 
-Coord = Annotated[Union[ArrayLike[TDims, TDtype], TDtype], FieldType.COORD]
+Coord = Annotated[Union[Collection[TDims, TDtype], TDtype], FieldType.COORD]
 """Type hint to define coordinate fields (``Coord[TDims, TDtype]``).
 
 Example:
@@ -189,7 +182,7 @@ class Image(AsDataArray):
 
 """
 
-Data = Annotated[Union[ArrayLike[TDims, TDtype], TDtype], FieldType.DATA]
+Data = Annotated[Union[Collection[TDims, TDtype], TDtype], FieldType.DATA]
 """Type hint to define data fields (``Coordof[TDims, TDtype]``).
 
 Examples:
@@ -267,7 +260,7 @@ def get_dims(type_: Any) -> Dims:
     args = get_args(type_)
     origin = get_origin(type_)
 
-    if origin is ArrayLike:
+    if origin is Collection:
         return get_dims(args[0])
 
     if origin is tuple or origin is Tuple:
@@ -298,7 +291,7 @@ def get_dtype(type_: Any) -> Dtype:
     args = get_args(type_)
     origin = get_origin(type_)
 
-    if origin is ArrayLike:
+    if origin is Collection:
         return get_dtype(args[1])
 
     if origin is Literal: