Repository: astropenguin/xarray-dataclasses (python)

Issue:
Add type hints for creating data/coord field from DataArray class

Pull request:
Add type hints for creating data/coord field from DataArray class
This PR closes #59.

Patch:
diff --git a/xarray_dataclasses/dataarray.py b/xarray_dataclasses/dataarray.py
index 6182a2a..b335ca3 100644
--- a/xarray_dataclasses/dataarray.py
+++ b/xarray_dataclasses/dataarray.py
@@ -52,7 +52,7 @@ def asdataarray(inst: Any, dataarray_factory: Any = xr.DataArray) -> Any:
     except AttributeError:
         pass
 
-    return parse(inst).to_dataarray(dataarray_factory)
+    return parse(inst).to_dataarray(dataarray_factory=dataarray_factory)
 
 
 def dataarrayclass(
diff --git a/xarray_dataclasses/dataset.py b/xarray_dataclasses/dataset.py
index d9eafba..a5716a6 100644
--- a/xarray_dataclasses/dataset.py
+++ b/xarray_dataclasses/dataset.py
@@ -47,7 +47,7 @@ def asdataset(inst: Any, dataset_factory: Any = xr.Dataset) -> Any:
     except AttributeError:
         pass
 
-    return parse(inst).to_dataset(dataset_factory)
+    return parse(inst).to_dataset(dataset_factory=dataset_factory)
 
 
 def datasetclass(
diff --git a/xarray_dataclasses/parser.py b/xarray_dataclasses/parser.py
index b2b851d..4339030 100644
--- a/xarray_dataclasses/parser.py
+++ b/xarray_dataclasses/parser.py
@@ -2,14 +2,14 @@ __all__ = ["parse"]
 
 
 # standard library
-from dataclasses import dataclass, Field
+from dataclasses import InitVar, dataclass, Field
 from typing import Any, List, Optional, Type, Union
 
 
 # third-party packages
 import numpy as np
 import xarray as xr
-from typing_extensions import get_args
+from typing_extensions import Protocol, TypedDict
 
 
 # submodules
@@ -25,78 +25,114 @@ from .typing import (
     TDataset,
     unannotate,
 )
+from .utils import resolve_class
 
 
 # type hints
 DataClassLike = Union[Type[DataClass], DataClass]
 Reference = Union[xr.DataArray, xr.Dataset]
+Types = TypedDict("Types", dims=Dims, dtype=Dtype)
 
 
 # dataclasses
 @dataclass(frozen=True)
-class DataArray:
-    """Parsed DataArray information."""
-
-    name: str
-    """Variable name for a DataArray."""
-    value: Any
-    """Value to be assigned to a DataArray."""
-    dims: Dims
-    """Dimensions of a DataArray."""
-    dtype: Dtype
-    """Data type of a DataArray."""
+class GeneralType:
+    """Representation for general-type variables."""
+
+    name: str  #: Name of a variable.
+    type: str  #: Type (full path) of a variable.
+    value: Any  #: Value to be assigned to a vabiable.
 
     @classmethod
-    def from_field(cls, field: Field[Any], value: Any) -> "DataArray":
+    def from_field(cls, field: Field[Any], value: Any) -> "GeneralType":
         """Create an instance from a dataclass field."""
-        t_dims, t_dtype = get_args(get_args(unannotate(field.type))[0])
-        return cls(field.name, value, get_dims(t_dims), get_dtype(t_dtype))
+        type = resolve_class(unannotate(field.type))
+        return cls(field.name, type, value)
+
+
+class DataArrayType(Protocol):
+    """Reperesentation for DataArray variables."""
+
+    name: str  #: Name of a variable
+    type: Union[Types, str]  #: Type (full path or dict) of a variable.
+    value: Any  #: Value to be assigned to a variable.
+
+    @classmethod
+    def from_field(cls, field: Field[Any], value: Any) -> "DataArrayType":
+        """Create an instance from a dataclass field."""
+        ...
 
     def __call__(self, reference: Optional[Reference] = None) -> xr.DataArray:
-        """Return a typed DataArray from the parsed information."""
-        return typed_dataarray(self.value, self.dims, self.dtype, reference)
+        """Create a typed DataArray from the representation."""
 
 
 @dataclass(frozen=True)
-class GeneralType:
-    """Parsed general-type information."""
+class DataType:
+    """Representation for DataArray variables with dims and dtypes."""
 
-    name: str
-    """Variable name for a type."""
-    value: Any
-    """Value to be assigned to a type."""
-    type: Type[Any]
-    """Type or type-hint object."""
+    name: str  #: Name of a variable
+    type: Types  #: Type (dict) of a variable.
+    value: Any  #: Value to be assigned to a variable.
 
     @classmethod
-    def from_field(cls, field: Field[Any], value: Any) -> "GeneralType":
+    def from_field(cls, field: Field[Any], value: Any) -> "DataType":
         """Create an instance from a dataclass field."""
-        return cls(field.name, value, unannotate(field.type))
+        type = unannotate(field.type).__args__[0]
 
-    def __call__(self) -> Any:
-        """Return the value (without casting)."""
-        return self.value
+        types: Types = {
+            "dims": get_dims(type.__args__[0]),
+            "dtype": get_dtype(type.__args__[1]),
+        }
+
+        return cls(field.name, types, value)
+
+    def __call__(self, reference: Optional[Reference] = None) -> xr.DataArray:
+        """Create a typed DataArray from the representation."""
+        dims, dtype = self.type["dims"], self.type["dtype"]
+        return typedarray(self.value, dims, dtype, reference)
 
 
 @dataclass(frozen=True)
-class DataStructure:
-    """Parsed dataclass information."""
-
-    attr: List[GeneralType]
-    """Parsed attr-type information."""
-    coord: List[DataArray]
-    """Parsed coord-type information."""
-    data: List[DataArray]
-    """Parsed data-type information."""
-    name: List[GeneralType]
-    """Parsed name-type information."""
+class ClassType:
+    """Representation for DataArray variables with dataclass."""
+
+    name: str  #: Name of a variable
+    type: str  #: Type (full path) of a variable.
+    value: Any  #: Value to be assigned to a variable.
+    dataclass: InitVar[Type[DataClass]]  #: Runtime dataclass of a variable.
+
+    def __post_init__(self, dataclass: InitVar[Type[DataClass]]) -> None:
+        super().__setattr__("dataclass", dataclass)
+
+    @classmethod
+    def from_field(cls, field: Field[Any], value: Any) -> "ClassType":
+        """Create an instance from a dataclass field."""
+        type = unannotate(field.type).__args__[0]
+        return cls(field.name, resolve_class(type), value, type)
+
+    def __call__(self, reference: Optional[Reference] = None) -> xr.DataArray:
+        """Create a typed DataArray from the representation."""
+        if isinstance(self.value, self.dataclass):
+            return parse(self.value).to_dataarray(reference)
+        else:
+            return parse(self.dataclass(self.value)).to_dataarray(reference)
+
+
+@dataclass(frozen=True)
+class Structure:
+    """Structure of a dataclass or its instance."""
+
+    attr: List[GeneralType]  #: Representations of attr-type variables.
+    coord: List[DataArrayType]  #: Representations of coord-type variables.
+    data: List[DataArrayType]  #: Representations of data-type variables.
+    name: List[GeneralType]  #: Representations of name-type variables.
 
     @classmethod
-    def from_dataclass(cls, dataclass: DataClassLike) -> "DataStructure":
-        """Create an instance from a dataclass."""
+    def from_dataclass(cls, dataclass: DataClassLike) -> "Structure":
+        """Create an instance from a dataclass or its instance."""
         attr: List[GeneralType] = []
-        coord: List[DataArray] = []
-        data: List[DataArray] = []
+        coord: List[DataArrayType] = []
+        data: List[DataArrayType] = []
         name: List[GeneralType] = []
 
         for field in dataclass.__dataclass_fields__.values():
@@ -105,9 +141,13 @@ class DataStructure:
             if FieldType.ATTR.annotates(field.type):
                 attr.append(GeneralType.from_field(field, value))
             elif FieldType.COORD.annotates(field.type):
-                coord.append(DataArray.from_field(field, value))
+                coord.append(DataType.from_field(field, value))
+            elif FieldType.COORDOF.annotates(field.type):
+                coord.append(ClassType.from_field(field, value))
             elif FieldType.DATA.annotates(field.type):
-                data.append(DataArray.from_field(field, value))
+                data.append(DataType.from_field(field, value))
+            elif FieldType.DATAOF.annotates(field.type):
+                data.append(ClassType.from_field(field, value))
             elif FieldType.NAME.annotates(field.type):
                 name.append(GeneralType.from_field(field, value))
 
@@ -115,70 +155,74 @@ class DataStructure:
 
     def to_dataarray(
         self,
+        reference: Optional[Reference] = None,
         dataarray_factory: Type[TDataArray] = xr.DataArray,
     ) -> TDataArray:
-        """Return a DataArray from the parsed information."""
-        return to_dataarray(self, dataarray_factory)
+        """Create a typed DataArray from the structure."""
+        return to_dataarray(self, reference, dataarray_factory)
 
     def to_dataset(
         self,
+        reference: Optional[Reference] = None,
         dataset_factory: Type[TDataset] = xr.Dataset,
     ) -> TDataset:
-        """Create a Dataset from the parsed information."""
-        return to_dataset(self, dataset_factory)
+        """Create a typed Dataset from the structure."""
+        return to_dataset(self, reference, dataset_factory)
 
 
 # runtime functions
-def parse(dataclass: DataClassLike) -> DataStructure:
-    """Create a data structure from a dataclass."""
-    return DataStructure.from_dataclass(dataclass)
+def parse(dataclass: DataClassLike) -> Structure:
+    """Create a structure from a dataclass."""
+    return Structure.from_dataclass(dataclass)
 
 
 def to_dataarray(
-    data_structure: DataStructure,
+    structure: Structure,
+    reference: Optional[Reference] = None,
     dataarray_factory: Type[TDataArray] = xr.DataArray,
 ) -> TDataArray:
-    """Create a DataArray from a parsed dataclass."""
-    dataarray = dataarray_factory(data_structure.data[0]())
+    """Create a typed DataArray from a structure."""
+    dataarray = dataarray_factory(structure.data[0](reference))
 
-    for coord in data_structure.coord:
+    for coord in structure.coord:
         dataarray.coords.update({coord.name: coord(dataarray)})
 
-    for attr in data_structure.attr:
-        dataarray.attrs.update({attr.name: attr()})
+    for attr in structure.attr:
+        dataarray.attrs.update({attr.name: attr.value})
 
-    for name in data_structure.name:
-        dataarray.name = name()
+    for name in structure.name:
+        dataarray.name = name.value
 
     return dataarray
 
 
 def to_dataset(
-    data_structure: DataStructure,
+    structure: Structure,
+    reference: Optional[Reference] = None,
     dataset_factory: Type[TDataset] = xr.Dataset,
 ) -> TDataset:
-    """Create a Dataset from a parsed dataclass."""
+    """Create a typed Dataset from a structure."""
     dataset = dataset_factory()
 
-    for data in data_structure.data:
-        dataset.update({data.name: data()})
+    for data in structure.data:
+        dataset.update({data.name: data(reference)})
 
-    for coord in data_structure.coord:
+    for coord in structure.coord:
         dataset.coords.update({coord.name: coord(dataset)})
 
-    for attr in data_structure.attr:
-        dataset.attrs.update({attr.name: attr()})
+    for attr in structure.attr:
+        dataset.attrs.update({attr.name: attr.value})
 
     return dataset
 
 
-def typed_dataarray(
+def typedarray(
     data: Any,
     dims: Dims,
     dtype: Dtype,
     reference: Optional[Reference] = None,
 ) -> xr.DataArray:
-    """Convert data to a DataArray with given dims and dtype."""
+    """Create a typed DataArray with given dims and dtype."""
     if not isinstance(data, ArrayLike):
         data = np.asarray(data)
 
@@ -195,10 +239,6 @@ def typed_dataarray(
     if reference is None:
         return dataarray
     else:
-        return dataarray.broadcast_like(subspace(reference, dims))
-
-
-def subspace(reference: Reference, dims: Dims) -> Reference:
-    """Return the subspace of a DataArray or a Dataset."""
-    diff_dims = set(reference.dims) - set(dims)
-    return reference.isel({dim: 0 for dim in diff_dims})
+        diff_dims = set(reference.dims) - set(dims)
+        subspace = reference.isel({dim: 0 for dim in diff_dims})
+        return dataarray.broadcast_like(subspace)
diff --git a/xarray_dataclasses/typing.py b/xarray_dataclasses/typing.py
index ed0c885..aaab21c 100644
--- a/xarray_dataclasses/typing.py
+++ b/xarray_dataclasses/typing.py
@@ -1,4 +1,4 @@
-__all__ = ["Attr", "Coord", "Data", "Name"]
+__all__ = ["Attr", "Coord", "Coordof", "Data", "Dataof", "Name"]
 
 
 # standard library
@@ -33,14 +33,12 @@ from typing_extensions import (
 class FieldType(Enum):
     """Annotation for xarray-related fields."""
 
-    ATTR = auto()
-    """Attribute field of DataArray or Dataset."""
-    COORD = auto()
-    """Coordinate field of DataArray or Dataset."""
-    DATA = auto()
-    """Data (variable) field of DataArray or Dataset."""
-    NAME = auto()
-    """Name field of DataArray."""
+    ATTR = auto()  #: Attribute field of DataArray or Dataset.
+    COORD = auto()  #: Coordinate field of DataArray or Dataset.
+    COORDOF = auto()  #: Coordinate field of DataArray or Dataset.
+    DATA = auto()  #: Data (variable) field of DataArray or Dataset.
+    DATAOF = auto()  #: Data (variable) field of DataArray or Dataset.
+    NAME = auto()  #: Name field of DataArray.
 
     def annotates(self, type_: Any) -> bool:
         """Check if type is annotated by the identifier."""
@@ -60,6 +58,7 @@ TDims = TypeVar("TDims", covariant=True)
 TDtype = TypeVar("TDtype", covariant=True)
 
 
+@runtime_checkable
 class DataClass(Protocol):
     """Type hint for dataclass objects."""
 
@@ -83,7 +82,7 @@ Examples:
 
         from dataclasses import dataclass
         from typing import Literal
-        from xarray_dataclasses import Data, Attr
+        from xarray_dataclasses import Attr, Data
 
 
         X = Literal["x"]
@@ -93,7 +92,7 @@ Examples:
         @dataclass
         class Image:
             data: Data[tuple[X, Y], float]
-            dpi: Attr[int] = 300
+            units: Attr[str] = "cd / m^2"
 
 """
 
@@ -105,7 +104,7 @@ Examples:
 
         from dataclasses import dataclass
         from typing import Literal
-        from xarray_dataclasses import Data, Coord
+        from xarray_dataclasses import Coord, Data
 
 
         X = Literal["x"]
@@ -115,12 +114,46 @@ Examples:
         @dataclass
         class Image:
             data: Data[tuple[X, Y], float]
-            weight: Coord[tuple[X, Y], float] = 1.0
             x: Coord[X, int] = 0
             y: Coord[Y, int] = 0
 
 """
 
+Coordof = Annotated[Union[T, ArrayLike[Any, Any], Any], FieldType.COORDOF]
+"""Type hint for a coordinate member of DataArray or Dataset.
+
+Unlike ``Coord``, it receives a dataclass that defines a DataArray class.
+
+Examples:
+    ::
+
+        from dataclasses import dataclass
+        from typing import Literal
+        from xarray_dataclasses import Data, Coordof
+
+
+        X = Literal["x"]
+        Y = Literal["y"]
+
+
+        @dataclass
+        class XAxis:
+            data: Data[X, int]
+
+
+        @dataclass
+        class YAxis:
+            data: Data[Y, int]
+
+
+        @dataclass
+        class Image:
+            data: Data[tuple[X, Y], float]
+            x: Coordof[XAxis] = 0
+            y: Coordof[YAxis] = 0
+
+"""
+
 Data = Annotated[Union[ArrayLike[TDims, TDtype], TDtype], FieldType.DATA]
 """Type hint for data of DataArray or variable of Dataset.
 
@@ -152,13 +185,55 @@ Examples:
 
 
         @dataclass
-        class Images:
+        class Image:
             red: Data[tuple[X, Y], float]
             green: Data[tuple[X, Y], float]
             blue: Data[tuple[X, Y], float]
 
 """
 
+Dataof = Annotated[Union[T, ArrayLike[Any, Any], Any], FieldType.DATAOF]
+"""Type hint for data of DataArray or variable of Dataset.
+
+Unlike ``Data``, it receives a dataclass that defines a DataArray class.
+
+Examples:
+    ::
+
+        from dataclasses import dataclass
+        from typing import Literal
+        from xarray_dataclasses import Coordof, Data, Dataof
+
+
+        X = Literal["x"]
+        Y = Literal["y"]
+
+
+        @dataclass
+        class XAxis:
+            data: Data[X, int]
+
+
+        @dataclass
+        class YAxis:
+            data: Data[Y, int]
+
+
+        @dataclass
+        class Image:
+            data: Data[tuple[X, Y], float]
+            x: Coordof[XAxis] = 0
+            y: Coordof[YAxis] = 0
+
+
+        @dataclass
+        class ColorImage:
+            red: Dataof[Image]
+            green: Dataof[Image]
+            red: Dataof[Image]
+
+"""
+
 Name = Annotated[T, FieldType.NAME]
 """Type hint for a name of DataArray.
 
@@ -177,7 +252,7 @@ Examples:
         @dataclass
         class Image:
             data: Data[tuple[X, Y], float]
-            name: Name[str] = "default"
+            name: Name[str] = "image"
 
 """
 
diff --git a/xarray_dataclasses/utils.py b/xarray_dataclasses/utils.py
index fb75e63..bd7f11c 100644
--- a/xarray_dataclasses/utils.py
+++ b/xarray_dataclasses/utils.py
@@ -2,11 +2,14 @@ __all__ = ["copy_class", "extend_class"]
 
 
 # standard library
-from typing import Any, Sequence, Type, TypeVar
+import re
+from pprint import pformat
+from typing import Any, Pattern, Sequence, Type, TypeVar
 
 
 # constants
 COPIED_CLASS: str = "__xrdc_copied_class__"
+CLASS_REPR: Pattern[str] = re.compile(r"^<class '(.+)'>$")
 
 
 # type hints
@@ -44,11 +47,12 @@ def extend_class(cls: Type[T], mixin: Type[Any]) -> Type[T]:
     return type(cls.__name__, bases, cls.__dict__.copy())
 
 
-def make_generic(cls: Type[T]) -> Type[T]:
-    """Make a class generic (only for type check)."""
-    try:
-        cls.__class_getitem__  # type: ignore
-    except AttributeError:
-        cls.__class_getitem__ = classmethod(GenericAlias)  # type: ignore
+def resolve_class(cls: Type[Any]) -> str:
+    """Return the prettified representation of a class."""
+    class_repr = pformat(cls)
+    match = CLASS_REPR.search(class_repr)
 
-    return cls
+    if match:
+        return match.group(1)
+    else:
+        return class_repr


Test patch:
diff --git a/tests/test_parser.py b/tests/test_parser.py
index 9397e06..61ba57b 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -9,7 +9,7 @@ from typing_extensions import Literal
 
 # submodules
 from xarray_dataclasses.parser import parse
-from xarray_dataclasses.typing import Attr, Coord, Data, Name
+from xarray_dataclasses.typing import Attr, Coord, Coordof, Data, Dataof, Name
 
 
 # type hints
@@ -19,40 +19,93 @@ Y = Literal["y"]
 
 # test datasets
 @dataclass
-class Image:
+class XAxis:
+    data: Data[X, int]
+    units: Attr[str] = "pixel"
+    name: Name[str] = "x axis"
+
+
+@dataclass
+class YAxis:
+    data: Data[Y, int]
+    units: Attr[str] = "pixel"
+    name: Name[str] = "y axis"
+
+
+@dataclass
+class Matrix:
     data: Data[Tuple[X, Y], float]
-    x: Coord[X, int] = 0
-    y: Coord[Y, int] = 0
-    dpi: Attr[int] = 100
-    name: Name[str] = "image"
+    mask: Coord[Tuple[X, Y], bool] = False
+    x: Coordof[XAxis] = 0
+    y: Coordof[YAxis] = 0
+
+
+@dataclass
+class Image:
+    red: Dataof[Matrix]
+    green: Dataof[Matrix]
+    blue: Dataof[Matrix]
 
 
-structure = parse(Image)
+xaxis_structure = parse(XAxis)
+yaxis_structure = parse(YAxis)
+matrix_structure = parse(Matrix)
+image_structure = parse(Image)
 
 
 # test functions
-def test_attr() -> None:
-    assert structure.attr[0].name == "dpi"
-    assert structure.attr[0].value == 100
-    assert structure.attr[0].type == int
+def test_xaxis_attr() -> None:
+    assert xaxis_structure.attr[0].name == "units"
+    assert xaxis_structure.attr[0].value == "pixel"
+    assert xaxis_structure.attr[0].type == "str"
+
+
+def test_xaxis_data() -> None:
+    assert xaxis_structure.data[0].name == "data"
+    assert xaxis_structure.data[0].type == {"dims": ("x",), "dtype": "int"}
+
+
+def test_xaxis_name() -> None:
+    assert xaxis_structure.name[0].name == "name"
+    assert xaxis_structure.name[0].value == "x axis"
+    assert xaxis_structure.name[0].type == "str"
+
+
+def test_yaxis_attr() -> None:
+    assert yaxis_structure.attr[0].name == "units"
+    assert yaxis_structure.attr[0].value == "pixel"
+    assert yaxis_structure.attr[0].type == "str"
+
+
+def test_yaxis_data() -> None:
+    assert yaxis_structure.data[0].name == "data"
+    assert yaxis_structure.data[0].type == {"dims": ("y",), "dtype": "int"}
+
+
+def test_yaxis_name() -> None:
+    assert yaxis_structure.name[0].name == "name"
+    assert yaxis_structure.name[0].value == "y axis"
+    assert yaxis_structure.name[0].type == "str"
 
 
-def test_coord() -> None:
-    assert structure.coord[0].name == "x"
-    assert structure.coord[0].dims == ("x",)
-    assert structure.coord[0].dtype == "int"
-    assert structure.coord[1].name == "y"
-    assert structure.coord[1].dims == ("y",)
-    assert structure.coord[1].dtype == "int"
+def test_matrix_coord() -> None:
+    assert matrix_structure.coord[0].name == "mask"
+    assert matrix_structure.coord[0].type == {"dims": ("x", "y"), "dtype": "bool"}
+    assert matrix_structure.coord[1].name == "x"
+    assert matrix_structure.coord[1].type == "tests.test_parser.XAxis"
+    assert matrix_structure.coord[2].name == "y"
+    assert matrix_structure.coord[2].type == "tests.test_parser.YAxis"
 
 
-def test_data() -> None:
-    assert structure.data[0].name == "data"
-    assert structure.data[0].dims == ("x", "y")
-    assert structure.data[0].dtype == "float"
+def test_matrix_data() -> None:
+    assert matrix_structure.data[0].name == "data"
+    assert matrix_structure.data[0].type == {"dims": ("x", "y"), "dtype": "float"}
 
 
-def test_name() -> None:
-    assert structure.name[0].name == "name"
-    assert structure.name[0].value == "image"
-    assert structure.name[0].type == str
+def test_image_data() -> None:
+    assert image_structure.data[0].name == "red"
+    assert image_structure.data[0].type == "tests.test_parser.Matrix"
+    assert image_structure.data[1].name == "green"
+    assert image_structure.data[1].type == "tests.test_parser.Matrix"
+    assert image_structure.data[2].name == "blue"
+    assert image_structure.data[2].type == "tests.test_parser.Matrix"