#141 Add labeled collection

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/typing.py

@@ -23,6 +23,7 @@
 from typing import (
     Any,
     ClassVar,
+    Collection,
     Dict,
     Hashable,
     Optional,
@@ -65,6 +66,18 @@
 Sizes = Dict[str, int]
 
 
+class Labeled(Protocol[TDims]):
+    """Type hint for labeled objects."""
+
+    pass
+
+
+class Collection(Labeled[TDims], Collection[TDtype], Protocol):
+    """Type hint for labeled collection objects."""
+
+    pass
+
+
 @runtime_checkable
 class ArrayLike(Protocol[TDims, TDtype]):
     """Type hint for array-like objects."""
@@ -138,7 +151,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 +202,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 +280,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 +311,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: