#158 Merge pull request from astropenguin/astropenguin/issue157

Update typing module

Author: astropenguin

poetry.lock

@@ -12,7 +12,6 @@ documentation = "https://astropenguin.github.io/xarray-dataclasses/"
 [tool.poetry.dependencies]
 python = ">=3.7.1, <3.11"
 morecopy = "^0.2"
-more-itertools = "^8.12"
 numpy = [
     { version = ">=1.15, <1.22", python = ">=3.7.1, <3.8" },
     { version = "^1.15", python = ">=3.8, <3.11" },
@@ -27,7 +26,7 @@ xarray = [
 black = "^22.3"
 ipython = [
     { version = "^7.32", python = ">=3.7.1, <3.8" },
-    { version = "^8.2", python = ">=3.8, <3.11" },
+    { version = "^8.4", python = ">=3.8, <3.11" },
 ]
 myst-parser = "^0.17"
 pydata-sphinx-theme = "^0.8"

pyproject.toml

@@ -1,87 +1,88 @@
 # standard library
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Tuple, Union
 
 
 # dependencies
+import numpy as np
 from pytest import mark
-from typing_extensions import Annotated, Literal
+from typing_extensions import Annotated as Ann
+from typing_extensions import Literal as L
 
 
 # submodules
 from xarray_dataclasses.typing import (
     Attr,
-    Collection,
     Coord,
     Data,
+    FType,
     Name,
     get_dims,
     get_dtype,
-    get_field_type,
-    get_repr_type,
+    get_ftype,
 )
 
 
-# type hints
-Int64 = Literal["int64"]
-NoneType = type(None)
-X = Literal["x"]
-Y = Literal["y"]
-
-
 # test datasets
 testdata_dims = [
-    (X, ("x",)),
-    (Tuple[()], ()),
-    (Tuple[X], ("x",)),
-    (Tuple[X, Y], ("x", "y")),
-    (Collection[X, Any], ("x",)),
-    (Collection[Tuple[()], Any], ()),
-    (Collection[Tuple[X], Any], ("x",)),
-    (Collection[Tuple[X, Y], Any], ("x", "y")),
+    (Coord[Tuple[()], Any], ()),
+    (Coord[L["x"], Any], ("x",)),
+    (Coord[Tuple[L["x"]], Any], ("x",)),
+    (Coord[Tuple[L["x"], L["y"]], Any], ("x", "y")),
+    (Data[Tuple[()], Any], ()),
+    (Data[L["x"], Any], ("x",)),
+    (Data[Tuple[L["x"]], Any], ("x",)),
+    (Data[Tuple[L["x"], L["y"]], Any], ("x", "y")),
+    (Ann[Coord[L["x"], Any], "coord"], ("x",)),
+    (Ann[Data[L["x"], Any], "data"], ("x",)),
+    (Union[Ann[Coord[L["x"], Any], "coord"], Ann[Any, "any"]], ("x",)),
+    (Union[Ann[Data[L["x"], Any], "data"], Ann[Any, "any"]], ("x",)),
 ]
 
 testdata_dtype = [
-    (Any, None),
-    (NoneType, None),
-    (Int64, "int64"),
-    (int, "int"),
-    (Collection[Any, Any], None),
-    (Collection[Any, NoneType], None),
-    (Collection[Any, Int64], "int64"),
-    (Collection[Any, int], "int"),
-]
-
-testdata_field_type = [
-    (Attr[Any], "attr"),
-    (Coord[Any, Any], "coord"),
-    (Data[Any, Any], "data"),
-    (Name[Any], "name"),
+    (Coord[Any, Any], None),
+    (Coord[Any, None], None),
+    (Coord[Any, int], np.dtype("i8")),
+    (Coord[Any, L["i8"]], np.dtype("i8")),
+    (Data[Any, Any], None),
+    (Data[Any, None], None),
+    (Data[Any, int], np.dtype("i8")),
+    (Data[Any, L["i8"]], np.dtype("i8")),
+    (Ann[Coord[Any, float], "coord"], np.dtype("f8")),
+    (Ann[Data[Any, float], "data"], np.dtype("f8")),
+    (Union[Ann[Coord[Any, float], "coord"], Ann[Any, "any"]], np.dtype("f8")),
+    (Union[Ann[Data[Any, float], "data"], Ann[Any, "any"]], np.dtype("f8")),
 ]
 
-testdata_repr_type = [
-    (int, int),
-    (Annotated[int, "annotation"], int),
-    (Union[int, float], int),
-    (Optional[int], int),
+testdata_ftype = [
+    (Attr[Any], FType.ATTR),
+    (Data[Any, Any], FType.DATA),
+    (Coord[Any, Any], FType.COORD),
+    (Name[Any], FType.NAME),
+    (Any, FType.OTHER),
+    (Ann[Attr[Any], "attr"], FType.ATTR),
+    (Ann[Data[Any, Any], "data"], FType.DATA),
+    (Ann[Coord[Any, Any], "coord"], FType.COORD),
+    (Ann[Name[Any], "name"], FType.NAME),
+    (Ann[Any, "other"], FType.OTHER),
+    (Union[Ann[Attr[Any], "attr"], Ann[Any, "any"]], FType.ATTR),
+    (Union[Ann[Data[Any, Any], "data"], Ann[Any, "any"]], FType.DATA),
+    (Union[Ann[Coord[Any, Any], "coord"], Ann[Any, "any"]], FType.COORD),
+    (Union[Ann[Name[Any], "name"], Ann[Any, "any"]], FType.NAME),
+    (Union[Ann[Any, "other"], Ann[Any, "any"]], FType.OTHER),
 ]
 
 
 # test functions
-@mark.parametrize("type_, dims", testdata_dims)
-def test_get_dims(type_: Any, dims: Any) -> None:
-    assert get_dims(type_) == dims
-
-
-@mark.parametrize("type_, dtype", testdata_dtype)
-def test_get_dtype(type_: Any, dtype: Any) -> None:
-    assert get_dtype(type_) == dtype
+@mark.parametrize("tp, dims", testdata_dims)
+def test_get_dims(tp: Any, dims: Any) -> None:
+    assert get_dims(tp) == dims
 
 
-@mark.parametrize("type_, field_type", testdata_field_type)
-def test_get_field_type(type_: Any, field_type: Any) -> None:
-    assert get_field_type(type_).value == field_type
+@mark.parametrize("tp, dtype", testdata_dtype)
+def test_get_dtype(tp: Any, dtype: Any) -> None:
+    assert get_dtype(tp) == dtype
 
 
-@mark.parametrize("type_, repr_type", testdata_repr_type)
-def test_get_repr_type(type_: Any, repr_type: Any) -> None:
-    assert get_repr_type(type_) == repr_type
+@mark.parametrize("tp, ftype", testdata_ftype)
+def test_get_ftype(tp: Any, ftype: Any) -> None:
+    assert get_ftype(tp) == ftype

tests/test_typing.py

@@ -5,16 +5,7 @@
 # standard library
 from functools import partial, wraps
 from types import MethodType
-from typing import (
-    Any,
-    Callable,
-    ClassVar,
-    Optional,
-    Type,
-    TypeVar,
-    Union,
-    overload,
-)
+from typing import Any, Callable, Optional, Type, TypeVar, Union, overload
 
 
 # dependencies
@@ -27,29 +18,25 @@
 # submodules
 from .datamodel import DataModel
 from .dataoptions import DataOptions
-from .typing import AnyArray, DataClass, DataClassFields, DataType, Order, Shape, Sizes
+from .typing import AnyArray, AnyXarray, DataClass, Order, Shape, Sizes
 
 
 # type hints
 PInit = ParamSpec("PInit")
 TDataArray = TypeVar("TDataArray", bound=xr.DataArray)
 
 
-class OptionedClass(Protocol[PInit, TDataArray]):
+class OptionedClass(DataClass[PInit], Protocol[PInit, TDataArray]):
     """Type hint for dataclass objects with options."""
 
-    def __init__(self, *args: PInit.args, **kwargs: PInit.kwargs) -> None:
-        ...
-
-    __dataclass_fields__: ClassVar[DataClassFields]
     __dataoptions__: DataOptions[TDataArray]
 
 
 # runtime functions
 @overload
 def asdataarray(
     dataclass: OptionedClass[PInit, TDataArray],
-    reference: Optional[DataType] = None,
+    reference: Optional[AnyXarray] = None,
     dataoptions: None = None,
 ) -> TDataArray:
     ...
@@ -58,7 +45,7 @@ def asdataarray(
 @overload
 def asdataarray(
     dataclass: DataClass[PInit],
-    reference: Optional[DataType] = None,
+    reference: Optional[AnyXarray] = None,
     dataoptions: None = None,
 ) -> xr.DataArray:
     ...
@@ -67,15 +54,15 @@ def asdataarray(
 @overload
 def asdataarray(
     dataclass: Any,
-    reference: Optional[DataType] = None,
+    reference: Optional[AnyXarray] = None,
     dataoptions: DataOptions[TDataArray] = DataOptions(xr.DataArray),
 ) -> TDataArray:
     ...
 
 
 def asdataarray(
     dataclass: Any,
-    reference: Optional[DataType] = None,
+    reference: Optional[AnyXarray] = None,
     dataoptions: Any = None,
 ) -> Any:
     """Create a DataArray object from a dataclass object.

xarray_dataclasses/dataarray.py

@@ -15,16 +15,17 @@
 
 # submodules
 from .typing import (
+    AnyDType,
     AnyField,
     DataClass,
-    DataType,
+    AnyXarray,
     Dims,
-    Dtype,
-    FieldType,
+    FType,
+    get_annotated,
+    get_dataclass,
     get_dims,
     get_dtype,
-    get_field_type,
-    get_repr_type,
+    get_ftype,
 )
 
 
@@ -93,7 +94,7 @@ class DataEntry:
     dims: Dims = cast(Dims, None)
     """Dimensions of the DataArray that the data is cast to."""
 
-    dtype: Dtype = cast(Dtype, None)
+    dtype: Optional[AnyDType] = None
     """Data type of the DataArray that the data is cast to."""
 
     base: Optional[Type[DataClass[Any]]] = None
@@ -119,7 +120,7 @@ def __post_init__(self) -> None:
         if model.names:
             setattr(self, "name", model.names[0].value)
 
-    def __call__(self, reference: Optional[DataType] = None) -> xr.DataArray:
+    def __call__(self, reference: Optional[AnyXarray] = None) -> xr.DataArray:
         """Create a DataArray object according to the entry."""
         from .dataarray import asdataarray
 
@@ -174,11 +175,11 @@ def from_dataclass(cls, dataclass: AnyDataClass[PInit]) -> "DataModel":
         eval_dataclass(dataclass)
 
         for field in dataclass.__dataclass_fields__.values():
-            try:
-                value = getattr(dataclass, field.name, MISSING)
-                model.entries[field.name] = get_entry(field, value)
-            except TypeError:
-                pass
+            value = getattr(dataclass, field.name, MISSING)
+            entry = get_entry(field, value)
+
+            if entry is not None:
+                model.entries[field.name] = entry
 
         return model
 
@@ -205,42 +206,41 @@ def eval_dataclass(dataclass: AnyDataClass[PInit]) -> None:
         field.type = types[field.name]
 
 
-def get_entry(field: AnyField, value: Any) -> AnyEntry:
+def get_entry(field: AnyField, value: Any) -> Optional[AnyEntry]:
     """Create an entry from a field and its value."""
-    field_type = get_field_type(field.type)
-    repr_type = get_repr_type(field.type)
+    ftype = get_ftype(field.type)
 
-    if field_type is FieldType.ATTR or field_type is FieldType.NAME:
+    if ftype is FType.ATTR or ftype is FType.NAME:
         return AttrEntry(
             name=field.name,
-            tag=field_type.value,
+            tag=ftype.value,
             value=value,
-            type=repr_type,
+            type=get_annotated(field.type),
         )
 
-    # hereafter field type is either COORD or DATA
-    if is_dataclass(repr_type):
-        return DataEntry(
-            name=field.name,
-            tag=field_type.value,
-            base=repr_type,
-            value=value,
-        )
-    else:
-        return DataEntry(
-            name=field.name,
-            tag=field_type.value,
-            dims=get_dims(repr_type),
-            dtype=get_dtype(repr_type),
-            value=value,
-        )
+    if ftype is FType.COORD or ftype is FType.DATA:
+        try:
+            return DataEntry(
+                name=field.name,
+                tag=ftype.value,
+                base=get_dataclass(field.type),
+                value=value,
+            )
+        except TypeError:
+            return DataEntry(
+                name=field.name,
+                tag=ftype.value,
+                dims=get_dims(field.type),
+                dtype=get_dtype(field.type),
+                value=value,
+            )
 
 
 def get_typedarray(
     data: Any,
     dims: Dims,
-    dtype: Dtype,
-    reference: Optional[DataType] = None,
+    dtype: Optional[AnyDType],
+    reference: Optional[AnyXarray] = None,
 ) -> xr.DataArray:
     """Create a DataArray object with given dims and dtype.
 

xarray_dataclasses/datamodel.py

@@ -8,17 +8,17 @@
 
 
 # submodules
-from .typing import DataType
+from .typing import AnyXarray
 
 
 # type hints
-TDataType = TypeVar("TDataType", bound=DataType)
+TAnyXarray = TypeVar("TAnyXarray", bound=AnyXarray)
 
 
 # dataclasses
 @dataclass(frozen=True)
-class DataOptions(Generic[TDataType]):
+class DataOptions(Generic[TAnyXarray]):
     """Options for DataArray or Dataset creation."""
 
-    factory: Callable[..., TDataType]
+    factory: Callable[..., TAnyXarray]
     """Factory function for DataArray or Dataset."""