#156 Update classmethod and runtime functions

Author: astropenguin

xarray_dataclasses/specs.py

@@ -1,21 +1,36 @@
-__all__ = ["DataSpec", "DataOptions"]
+__all__ = ["DataOptions", "DataSpec"]
 
 
 # standard library
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, fields
+from functools import lru_cache
 from typing import Any, Dict, Generic, Hashable, Optional, Type, TypeVar
 
 
 # dependencies
-from typing_extensions import Literal, TypeAlias
+from typing_extensions import Literal, TypeAlias, get_type_hints
 
 
 # submodules
-from .typing import AnyDType, AnyXarray, DataClass, Dims
+from .typing import (
+    AnyDType,
+    AnyField,
+    AnyXarray,
+    DataClass,
+    Dims,
+    Role,
+    get_annotated,
+    get_dataclass,
+    get_dims,
+    get_dtype,
+    get_name,
+    get_role,
+)
 
 
 # type hints
 AnySpec: TypeAlias = "ArraySpec | ScalarSpec"
+TDataClass = TypeVar("TDataClass", bound=DataClass[...])
 TReturn = TypeVar("TReturn", AnyXarray, None)
 
 
@@ -33,14 +48,31 @@ class ArraySpec:
     default: Any
     """Default value of the array."""
 
-    dims: Dims
+    dims: Dims = ()
     """Dimensions of the array."""
 
-    type: Optional[AnyDType]
+    type: Optional[AnyDType] = None
     """Data type of the array."""
 
     origin: Optional[Type[DataClass[Any]]] = None
-    """Dataclass of dims and type origins."""
+    """Dataclass as origins of name, dims, and type."""
+
+    def __post_init__(self) -> None:
+        """Update name, dims, and type if origin exists."""
+        if self.origin is None:
+            return
+
+        dataspec = DataSpec.from_dataclass(self.origin)
+        setattr = object.__setattr__
+
+        for spec in dataspec.specs.of_data.values():
+            setattr(self, "dims", spec.dims)
+            setattr(self, "type", spec.type)
+            break
+
+        for spec in dataspec.specs.of_name.values():
+            setattr(self, "name", spec.default)
+            break
 
 
 @dataclass(frozen=True)
@@ -101,3 +133,63 @@ class DataSpec:
 
     options: DataOptions[Any] = DataOptions(type(None))
     """Options for xarray data creation."""
+
+    @classmethod
+    def from_dataclass(cls, dataclass: Type[DataClass[...]]) -> "DataSpec":
+        """Create a data specification from a dataclass."""
+        specs = SpecDict()
+
+        for field in fields(eval_fields(dataclass)):
+            spec = get_spec(field)
+
+            if spec is not None:
+                specs[field.name] = spec
+
+        try:
+            return cls(specs, dataclass.__dataoptions__)  # type: ignore
+        except AttributeError:
+            return cls(specs)
+
+
+# runtime functions
+@lru_cache(maxsize=None)
+def eval_fields(dataclass: Type[TDataClass]) -> Type[TDataClass]:
+    """Evaluate field types of a dataclass."""
+    types = get_type_hints(dataclass, include_extras=True)
+
+    for field in fields(dataclass):
+        field.type = types[field.name]
+
+    return dataclass
+
+
+@lru_cache(maxsize=None)
+def get_spec(field: AnyField) -> Optional[AnySpec]:
+    """Convert a dataclass field to a specification."""
+    name = get_name(field.type, field.name)
+    role = get_role(field.type)
+
+    if role is Role.DATA or role is Role.COORD:
+        try:
+            return ArraySpec(
+                name=name,
+                role=role.value,
+                default=field.default,
+                origin=get_dataclass(field.type),
+            )
+        except TypeError:
+            return ArraySpec(
+                name=name,
+                role=role.value,
+                default=field.default,
+                dims=get_dims(field.type),
+                type=get_dtype(field.type),
+            )
+
+    if role is Role.ATTR or role is Role.NAME:
+        return ScalarSpec(
+            name=name,
+            role=role.value,
+            default=field.default,
+            type=get_annotated(field.type),
+        )