#237 Add helper dataclass converters

Author: astropenguin

xarray_dataclasses/api.py

@@ -2,12 +2,31 @@
 
 
 # standard library
-from typing import Any, Callable, overload
-
+from dataclasses import replace
+from typing import Any, ForwardRef, Literal, Optional, overload
 
 # dependencies
-from xarray import DataArray, Dataset
-from .typing import DataClass, DataClassOf, PAny, TDataArray, TDataset, TXarray
+from dataspecs import ID, ROOT, Spec, Specs
+from numpy import asarray, array
+from typing_extensions import get_args, get_origin
+from xarray import DataArray, Dataset, Variable
+from .typing import (
+    DataClass,
+    DataClassOf,
+    Factory,
+    HashDict,
+    PAny,
+    TAny,
+    TDataArray,
+    TDataset,
+    TXarray,
+    Tag,
+)
+
+
+# type hints
+Attrs = HashDict[Any]
+Vars = HashDict[Variable]
 
 
 @overload
@@ -24,7 +43,7 @@ def asdataarray(
     obj: DataClass[PAny],
     /,
     *,
-    factory: Callable[..., TDataArray],
+    factory: Factory[TDataArray],
 ) -> TDataArray: ...
 
 
@@ -56,7 +75,7 @@ def asdataset(
     obj: DataClass[PAny],
     /,
     *,
-    factory: Callable[..., TDataset],
+    factory: Factory[TDataset],
 ) -> TDataset: ...
 
 
@@ -88,10 +107,84 @@ def asxarray(
     obj: DataClass[PAny],
     /,
     *,
-    factory: Callable[..., TXarray],
+    factory: Factory[TXarray],
 ) -> TXarray: ...
 
 
 def asxarray(obj: Any, /, *, factory: Any = None) -> Any:
     """Create a DataArray/set object from a dataclass object."""
     ...
+
+
+def get_attrs(specs: Specs[Spec[Any]], /, *, at: ID = ROOT) -> Attrs:
+    """Create attributes from data specs."""
+    attrs: Attrs = {}
+
+    for spec in specs[at.children][Tag.ATTR]:
+        options = specs[spec.id.children]
+        factory = maybe(options[Tag.FACTORY].unique).data or identity
+        name = maybe(options[Tag.NAME].unique).data or spec.id.name
+
+        if Tag.MULTIPLE not in spec.tags:
+            spec = replace(spec, data={name: spec.data})
+
+        for name, data in spec[HashDict[Any]].data.items():
+            attrs[name] = factory(data)
+
+    return attrs
+
+
+def get_vars(specs: Specs[Spec[Any]], of: Tag, /, *, at: ID = ROOT) -> Vars:
+    """Create variables of given tag from data specs."""
+    vars: Vars = {}
+
+    for spec in specs[at.children][of]:
+        options = specs[spec.id.children]
+        attrs = get_attrs(specs, at=spec.id)
+        factory = maybe(options[Tag.FACTORY].unique).data or Variable
+        name = maybe(options[Tag.NAME].unique).data or spec.id.name
+
+        if (type_ := maybe(options[Tag.DIMS].unique).type) is None:
+            raise RuntimeError("Could not find any data spec for dims.")
+        elif get_origin(type_) is tuple:
+            dims = tuple(str(unwrap(arg)) for arg in get_args(type_))
+        else:
+            dims = (str(unwrap(type_)),)
+
+        if (type_ := maybe(options[Tag.DTYPE].unique).type) is None:
+            raise RuntimeError("Could not find any data spec for dims.")
+        elif type_ is type(None) or type_ is Any:
+            dtype = None
+        else:
+            dtype = unwrap(type_)
+
+        if Tag.MULTIPLE not in spec.tags:
+            spec = replace(spec, data={name: spec.data})
+
+        for name, data in spec[HashDict[Any]].data.items():
+            if not (data := asarray(data, dtype)).ndim:
+                data = array(data, ndmin=len(dims))
+
+            vars[name] = factory(attrs=attrs, data=data, dims=dims)
+
+    return vars
+
+
+def identity(obj: TAny, /) -> TAny:
+    """Identity function used for the default factory."""
+    return obj
+
+
+def maybe(obj: Optional[Spec[Any]], /) -> Spec[Any]:
+    """Return a dummy (``None``-filled) data spec if an object is not one."""
+    return Spec(ROOT, (), None, None) if obj is None else obj
+
+
+def unwrap(obj: Any, /) -> Any:
+    """Unwrap if an object is a literal or a forward reference."""
+    if get_origin(obj) is Literal:
+        return args[0] if len(args := get_args(obj)) == 1 else obj
+    elif isinstance(obj, ForwardRef):
+        return obj.__forward_arg__
+    else:
+        return obj

xarray_dataclasses/typing.py

@@ -41,6 +41,7 @@
 TDtype = TypeVar("TDtype", covariant=True)
 THashable = TypeVar("THashable", bound=Hashable)
 TXarray = TypeVar("TXarray", covariant=True, bound="Xarray")
+HashDict = dict[Hashable, TAny]
 Xarray = Union[DataArray, Dataset]
 
 
@@ -103,19 +104,19 @@ class Tag(TagBase):
 Attr = Annotated[TAny, Tag.ATTR]
 """Type alias for an attribute of DataArray/set."""
 
-Attrs = Annotated[dict[str, TAny], Tag.ATTR, Tag.MULTIPLE]
+Attrs = Annotated[HashDict[TAny], Tag.ATTR, Tag.MULTIPLE]
 """Type alias for attributes of DataArray/set."""
 
 Coord = Annotated[Arrayable[TDims, TDtype], Tag.COORD]
 """Type alias for a coordinate of DataArray/set."""
 
-Coords = Annotated[dict[str, Arrayable[TDims, TDtype]], Tag.COORD, Tag.MULTIPLE]
+Coords = Annotated[HashDict[Arrayable[TDims, TDtype]], Tag.COORD, Tag.MULTIPLE]
 """Type alias for coordinates of DataArray/set."""
 
 Data = Annotated[Arrayable[TDims, TDtype], Tag.DATA]
 """Type alias for a data object of DataArray/set."""
 
-DataVars = Annotated[dict[str, Arrayable[TDims, TDtype]], Tag.DATA, Tag.MULTIPLE]
+DataVars = Annotated[HashDict[Arrayable[TDims, TDtype]], Tag.DATA, Tag.MULTIPLE]
 """Type alias for data objects of DataArray/set."""
 
 Factory = Annotated[Callable[..., TXarray], Tag.FACTORY]