#137 Add runtime codes for DataModel

Author: astropenguin

xarray_dataclasses/datamodel.py

@@ -2,7 +2,7 @@
 
 
 # standard library
-from dataclasses import dataclass, field, is_dataclass
+from dataclasses import Field, dataclass, field, is_dataclass
 from typing import Any, Dict, Hashable, List, Optional, Tuple, Type, Union, cast
 
 
@@ -13,7 +13,18 @@
 
 
 # submodules
-from .typing import ArrayLike, DataClass, DataType, Dims, Dtype
+from .typing import (
+    ArrayLike,
+    DataClass,
+    DataType,
+    Dims,
+    Dtype,
+    FieldType,
+    get_dims,
+    get_dtype,
+    get_field_type,
+    get_repr_type,
+)
 
 
 # type hints
@@ -63,7 +74,10 @@ class AttrEntry:
 
     def __call__(self) -> Any:
         """Create an object according to the entry."""
-        ...
+        if self.value is MISSING:
+            raise ValueError("Value is missing.")
+
+        return self.value
 
 
 @dataclass(frozen=True)
@@ -82,7 +96,7 @@ class DataEntry:
     dtype: Dtype = cast(Dtype, None)
     """Data type of the DataArray that the data is cast to."""
 
-    base: Optional[Type[Any]] = None
+    base: Optional[Type[DataClass[Any]]] = None
     """Base dataclass that converts the data to a DataArray."""
 
     value: Any = MISSING
@@ -91,9 +105,34 @@ class DataEntry:
     cast: bool = True
     """Whether the value is cast to the data type."""
 
+    def __post_init__(self) -> None:
+        """Update the entry if a base dataclass exists."""
+        if self.base is None:
+            return
+
+        model = DataModel.from_dataclass(self.base)
+
+        setattr = object.__setattr__
+        setattr(self, "dims", model.data_vars[0].dims)
+        setattr(self, "dtype", model.data_vars[0].dtype)
+
+        if model.names:
+            setattr(self, "name", model.names[0].value)
+
     def __call__(self, reference: Optional[DataType] = None) -> xr.DataArray:
         """Create a DataArray object according to the entry."""
-        ...
+        from .dataarray import asdataarray
+
+        if self.value is MISSING:
+            raise ValueError("Value is missing.")
+
+        if self.base is None:
+            return get_typedarray(self.value, self.dims, self.dtype, reference)
+
+        if is_dataclass(self.value):
+            return asdataarray(self.value, reference)
+        else:
+            return asdataarray(self.base(self.value), reference)
 
 
 @dataclass(frozen=True)
@@ -106,32 +145,42 @@ class DataModel:
     @property
     def attrs(self) -> List[AttrEntry]:
         """Return a list of attribute entries."""
-        ...
+        return [v for v in self.entries.values() if v.tag == "attr"]
 
     @property
     def coords(self) -> List[DataEntry]:
         """Return a list of coordinate entries."""
-        ...
+        return [v for v in self.entries.values() if v.tag == "coord"]
 
     @property
     def data_vars(self) -> List[DataEntry]:
         """Return a list of data variable entries."""
-        ...
+        return [v for v in self.entries.values() if v.tag == "data"]
 
     @property
     def data_vars_items(self) -> List[Tuple[str, DataEntry]]:
         """Return a list of data variable entries with keys."""
-        ...
+        return [(k, v) for k, v in self.entries.items() if v.tag == "data"]
 
     @property
     def names(self) -> List[AttrEntry]:
         """Return a list of name entries."""
-        ...
+        return [v for v in self.entries.values() if v.tag == "name"]
 
     @classmethod
     def from_dataclass(cls, dataclass: AnyDataClass[P]) -> "DataModel":
         """Create a data model from a dataclass or its object."""
-        ...
+        model = cls()
+        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
+
+        return model
 
 
 # runtime functions
@@ -156,7 +205,38 @@ def eval_dataclass(dataclass: AnyDataClass[P]) -> None:
         field.type = types[field.name]
 
 
-def typedarray(
+def get_entry(field: Field[Any], value: Any) -> AnyEntry:
+    """Create an entry from a field and its value."""
+    field_type = get_field_type(field.type)
+    repr_type = get_repr_type(field.type)
+
+    if field_type is FieldType.ATTR or field_type is FieldType.NAME:
+        return AttrEntry(
+            name=field.name,
+            tag=field_type.value,
+            value=value,
+            type=repr_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,
+        )
+
+
+def get_typedarray(
     data: Any,
     dims: Dims,
     dtype: Dtype,

xarray_dataclasses/typing.py

@@ -48,33 +48,6 @@
 )
 
 
-# constants
-class FieldType(Enum):
-    """Annotation of xarray-related field hints."""
-
-    ATTR = "attr"
-    """Annotation of attribute field hints."""
-
-    COORD = "coord"
-    """Annotation of coordinate field hints."""
-
-    COORDOF = "coordof"
-    """Annotation of coordinate field hints."""
-
-    DATA = "data"
-    """Annotation of data (variable) field hints."""
-
-    DATAOF = "dataof"
-    """Annotation of data (variable) field hints."""
-
-    NAME = "name"
-    """Annotation of name field hints."""
-
-    def annotates(self, hint: Any) -> bool:
-        """Check if a field hint is annotated."""
-        return self in get_args(hint)[1:]
-
-
 # type hints
 P = ParamSpec("P")
 T = TypeVar("T")
@@ -120,6 +93,28 @@ def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
     __dataclass_fields__: ClassVar[DataClassFields]
 
 
+# constants
+class FieldType(Enum):
+    """Annotation of xarray-related field hints."""
+
+    ATTR = "attr"
+    """Annotation of attribute field hints."""
+
+    COORD = "coord"
+    """Annotation of coordinate field hints."""
+
+    DATA = "data"
+    """Annotation of data (variable) field hints."""
+
+    NAME = "name"
+    """Annotation of name field hints."""
+
+    def annotates(self, hint: Any) -> bool:
+        """Check if a field hint is annotated."""
+        return self in get_args(hint)[1:]
+
+
+# public type hints
 Attr = Annotated[T, FieldType.ATTR]
 """Type hint to define attribute fields (``Attr[T]``).
 
@@ -162,7 +157,7 @@ class Image(AsDataArray):
 
 """
 
-Coordof = Annotated[Union[TDataClass, Any], FieldType.COORDOF]
+Coordof = Annotated[Union[TDataClass, Any], FieldType.COORD]
 """Type hint to define coordinate fields (``Coordof[TDataClass]``).
 
 Unlike ``Coord``, it specifies a dataclass that defines a DataArray class.
@@ -215,7 +210,7 @@ class ColorImage(AsDataset):
 
 """
 
-Dataof = Annotated[Union[TDataClass, Any], FieldType.DATAOF]
+Dataof = Annotated[Union[TDataClass, Any], FieldType.DATA]
 """Type hint to define data fields (``Coordof[TDataClass]``).
 
 Unlike ``Data``, it specifies a dataclass that defines a DataArray class.