Repository: astropenguin/xarray-dataclasses

Issue:
Update NumPy-like zeros, ones, ... to be compatible with sizes

Discussion:

Patch:
diff --git a/xarray_dataclasses/dataarray.py b/xarray_dataclasses/dataarray.py
index 3db7135..6cce9a1 100644
--- a/xarray_dataclasses/dataarray.py
+++ b/xarray_dataclasses/dataarray.py
@@ -24,7 +24,7 @@ P = ParamSpec("P")
 TDataArray = TypeVar("TDataArray", bound=xr.DataArray)
 TDataArray_ = TypeVar("TDataArray_", bound=xr.DataArray, contravariant=True)
 Order = Literal["C", "F"]
-Shape = Union[Sequence[int], int]
+Shape = Union[Dict[str, int], Sequence[int], int]
 
 
 class DataClass(Protocol[P]):
@@ -161,7 +161,13 @@ class AsDataArray:
             DataArray object filled without initializing data.
 
         """
-        name = DataModel.from_dataclass(cls).data[0].name
+        model = DataModel.from_dataclass(cls)
+        name = model.data[0].name
+        dims = model.data[0].type["dims"]
+
+        if isinstance(shape, dict):
+            shape = tuple(shape[dim] for dim in dims)
+
         data = np.empty(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
 
@@ -184,7 +190,13 @@ class AsDataArray:
             DataArray object filled with zeros.
 
         """
-        name = DataModel.from_dataclass(cls).data[0].name
+        model = DataModel.from_dataclass(cls)
+        name = model.data[0].name
+        dims = model.data[0].type["dims"]
+
+        if isinstance(shape, dict):
+            shape = tuple(shape[dim] for dim in dims)
+
         data = np.zeros(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
 
@@ -207,7 +219,13 @@ class AsDataArray:
             DataArray object filled with ones.
 
         """
-        name = DataModel.from_dataclass(cls).data[0].name
+        model = DataModel.from_dataclass(cls)
+        name = model.data[0].name
+        dims = model.data[0].type["dims"]
+
+        if isinstance(shape, dict):
+            shape = tuple(shape[dim] for dim in dims)
+
         data = np.ones(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
 
@@ -232,6 +250,12 @@ class AsDataArray:
             DataArray object filled with given value.
 
         """
-        name = DataModel.from_dataclass(cls).data[0].name
+        model = DataModel.from_dataclass(cls)
+        name = model.data[0].name
+        dims = model.data[0].type["dims"]
+
+        if isinstance(shape, dict):
+            shape = tuple(shape[dim] for dim in dims)
+
         data = np.full(shape, fill_value, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
diff --git a/xarray_dataclasses/datamodel.py b/xarray_dataclasses/datamodel.py
index 7d49338..a865899 100644
--- a/xarray_dataclasses/datamodel.py
+++ b/xarray_dataclasses/datamodel.py
@@ -2,8 +2,8 @@ __all__ = ["DataModel"]
 
 
 # standard library
-from dataclasses import Field, InitVar, dataclass, field, is_dataclass
-from typing import Any, List, Type, Union, cast
+from dataclasses import Field, dataclass, field, is_dataclass
+from typing import Any, List, Optional, Type, Union, cast
 
 
 # dependencies
@@ -22,91 +22,99 @@ from .typing import (
     FieldType,
     get_dims,
     get_dtype,
-    get_first,
-    get_repr,
+    get_inner,
+    unannotate,
 )
 
 
 # type hints
+DataType = TypedDict("DataType", dims=Dims, dtype=Dtype)
 Reference = Union[xr.DataArray, xr.Dataset, None]
-DataTypes = TypedDict("DataTypes", dims=Dims, dtype=Dtype)
 
 
 # field models
-@dataclass
+@dataclass(frozen=True)
 class Data:
-    """Model for the coord or data fields."""
+    """Field model for data-related fields."""
 
     name: str
-    type: DataTypes
-    value: Any
+    """Name of the field."""
 
-    def __call__(self, reference: Reference = None) -> xr.DataArray:
-        """Create a DataArray object from the value and a reference."""
-        return typedarray(
-            self.value,
-            self.type["dims"],
-            self.type["dtype"],
-            reference,
-        )
-
-    @classmethod
-    def from_field(cls, field: Field[Any], value: Any) -> "Data":
-        """Create a field model from a dataclass field and a value."""
-        return cls(
-            field.name,
-            {
-                "dims": get_dims(field.type),
-                "dtype": get_dtype(field.type),
-            },
-            value,
-        )
-
-
-@dataclass
-class Dataof:
-    """Model for the coordof or dataof fields."""
-
-    name: str
-    type: str
     value: Any
-    dataclass: InitVar[Type[DataClass]]
+    """Value assigned to the field."""
+
+    type: DataType
+    """Type (dims and dtype) of the field."""
 
-    def __post_init__(self, dataclass: Type[DataClass]) -> None:
-        self._dataclass = dataclass
+    factory: Optional[Type[DataClass]] = None
+    """Factory dataclass to create a DataArray object."""
 
     def __call__(self, reference: Reference = None) -> xr.DataArray:
         """Create a DataArray object from the value and a reference."""
         from .dataarray import asdataarray
 
+        if self.factory is None:
+            return typedarray(
+                self.value,
+                self.type["dims"],
+                self.type["dtype"],
+                reference,
+            )
+
         if is_dataclass(self.value):
             return asdataarray(self.value, reference)
         else:
-            return asdataarray(self._dataclass(self.value), reference)
+            return asdataarray(self.factory(self.value), reference)
 
     @classmethod
-    def from_field(cls, field: Field[Any], value: Any) -> "Dataof":
+    def from_field(cls, field: Field[Any], value: Any, of: bool) -> "Data":
         """Create a field model from a dataclass field and a value."""
-        dataclass = get_first(field.type)
-        return cls(field.name, get_repr(dataclass), value, dataclass)
+        hint = unannotate(field.type)
 
+        if of:
+            dataclass = get_inner(hint, 0)
+            data = DataModel.from_dataclass(dataclass).data[0]
+            return cls(field.name, value, data.type, dataclass)
+        else:
+            return cls(
+                field.name,
+                value,
+                {"dims": get_dims(hint), "dtype": get_dtype(hint)},
+            )
 
-@dataclass
+
+@dataclass(frozen=True)
 class General:
-    """Model for the attribute or name fields."""
+    """Field model for general fields."""
 
     name: str
-    type: str
+    """Name of the field."""
+
     value: Any
+    """Value assigned to the field."""
+
+    type: str
+    """Type of the field."""
+
+    factory: Optional[Type[Any]] = None
+    """Factory function to create an object."""
 
     def __call__(self) -> Any:
-        """Just return the value."""
-        return self.value
+        """Create an object from the value."""
+        if self.factory is None:
+            return self.value
+        else:
+            return self.factory(self.value)
 
     @classmethod
     def from_field(cls, field: Field[Any], value: Any) -> "General":
         """Create a field model from a dataclass field and a value."""
-        return cls(field.name, get_repr(field.type), value)
+        hint = unannotate(field.type)
+
+        try:
+            return cls(field.name, value, f"{hint.__module__}.{hint.__qualname__}")
+        except AttributeError:
+            return cls(field.name, value, repr(hint))
 
 
 # data models
@@ -117,10 +125,10 @@ class DataModel:
     attr: List[General] = field(default_factory=list)
     """Model of the attribute fields."""
 
-    coord: List[Union[Data, Dataof]] = field(default_factory=list)
+    coord: List[Data] = field(default_factory=list)
     """Model of the coordinate fields."""
 
-    data: List[Union[Data, Dataof]] = field(default_factory=list)
+    data: List[Data] = field(default_factory=list)
     """Model of the data fields."""
 
     name: List[General] = field(default_factory=list)
@@ -138,13 +146,13 @@ class DataModel:
             if FieldType.ATTR.annotates(field_.type):
                 model.attr.append(General.from_field(field_, value))
             elif FieldType.COORD.annotates(field_.type):
-                model.coord.append(Data.from_field(field_, value))
+                model.coord.append(Data.from_field(field_, value, False))
             elif FieldType.COORDOF.annotates(field_.type):
-                model.coord.append(Dataof.from_field(field_, value))
+                model.coord.append(Data.from_field(field_, value, True))
             elif FieldType.DATA.annotates(field_.type):
-                model.data.append(Data.from_field(field_, value))
+                model.data.append(Data.from_field(field_, value, False))
             elif FieldType.DATAOF.annotates(field_.type):
-                model.data.append(Dataof.from_field(field_, value))
+                model.data.append(Data.from_field(field_, value, True))
             elif FieldType.NAME.annotates(field_.type):
                 model.name.append(General.from_field(field_, value))
 
diff --git a/xarray_dataclasses/typing.py b/xarray_dataclasses/typing.py
index 55b1c4f..05c571c 100644
--- a/xarray_dataclasses/typing.py
+++ b/xarray_dataclasses/typing.py
@@ -14,7 +14,6 @@ __all__ = ["Attr", "Coord", "Coordof", "Data", "Dataof", "Name"]
 
 
 # standard library
-import re
 from dataclasses import Field
 from enum import auto, Enum
 from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union
@@ -23,11 +22,11 @@ from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union
 # dependencies
 from typing_extensions import (
     Annotated,
+    Literal,
+    Protocol,
     get_args,
     get_origin,
     get_type_hints,
-    Literal,
-    Protocol,
     runtime_checkable,
 )
 
@@ -235,10 +234,10 @@ Example:
 # runtime functions
 def get_dims(hint: Any) -> Dims:
     """Return dims parsed from a type hint."""
-    t_dims = get_args(get_first(hint))[0]
+    t_dims = get_inner(hint, 0, 0)
 
     if is_str_literal(t_dims):
-        return (get_first(t_dims),)
+        return (get_inner(t_dims, 0),)
 
     args = get_args(t_dims)
 
@@ -246,14 +245,14 @@ def get_dims(hint: Any) -> Dims:
         return ()
 
     if all(map(is_str_literal, args)):
-        return tuple(map(get_first, args))
+        return tuple(map(get_inner, args, [0] * len(args)))
 
     raise ValueError(f"Could not parse dims from {hint!r}.")
 
 
 def get_dtype(hint: Any) -> Dtype:
     """Return dtype parsed from a type hint."""
-    t_dtype = get_args(get_first(hint))[1]
+    t_dtype = get_inner(hint, 0, 1)
 
     if t_dtype is Any:
         return None
@@ -265,25 +264,18 @@ def get_dtype(hint: Any) -> Dtype:
         return t_dtype.__name__
 
     if is_str_literal(t_dtype):
-        return get_first(t_dtype)
+        return get_inner(t_dtype, 0)
 
     raise ValueError(f"Could not parse dtype from {hint!r}.")
 
 
-def get_first(hint: Any) -> Any:
-    """Return the first argument in a type hint."""
-    return get_args(unannotate(hint))[0]
-
-
-def get_repr(hint: Any) -> str:
-    """Return the representation of a type hint."""
-    hint_repr = repr(unannotate(hint))
-    match = re.search(r"^<class '(.+)'>$", hint_repr)
+def get_inner(hint: Any, *indexes: int) -> Any:
+    """Return an inner type hint by indexes."""
+    if not indexes:
+        return hint
 
-    if match:
-        return match.group(1)
-    else:
-        return hint_repr
+    index, indexes = indexes[0], indexes[1:]
+    return get_inner(get_args(hint)[index], *indexes)
 
 
 def is_str_literal(hint: Any) -> bool:


Test patch:
diff --git a/tests/test_datamodel.py b/tests/test_datamodel.py
index 28d334b..3a3f004 100644
--- a/tests/test_datamodel.py
+++ b/tests/test_datamodel.py
@@ -57,44 +57,51 @@ color_model = DataModel.from_dataclass(ColorImage)
 def test_xaxis_attr() -> None:
     assert xaxis_model.attr[0].name == "units"
     assert xaxis_model.attr[0].value == "pixel"
-    assert xaxis_model.attr[0].type == "str"
+    assert xaxis_model.attr[0].type == "builtins.str"
 
 
 def test_xaxis_data() -> None:
     assert xaxis_model.data[0].name == "data"
     assert xaxis_model.data[0].type == {"dims": ("x",), "dtype": "int"}
+    assert xaxis_model.data[0].factory is None
 
 
 def test_xaxis_name() -> None:
     assert xaxis_model.name[0].name == "name"
     assert xaxis_model.name[0].value == "x axis"
-    assert xaxis_model.name[0].type == "str"
+    assert xaxis_model.name[0].type == "builtins.str"
 
 
 def test_yaxis_attr() -> None:
     assert yaxis_model.attr[0].name == "units"
     assert yaxis_model.attr[0].value == "pixel"
-    assert yaxis_model.attr[0].type == "str"
+    assert yaxis_model.attr[0].type == "builtins.str"
 
 
 def test_yaxis_data() -> None:
     assert yaxis_model.data[0].name == "data"
     assert yaxis_model.data[0].type == {"dims": ("y",), "dtype": "int"}
+    assert yaxis_model.data[0].factory is None
 
 
 def test_yaxis_name() -> None:
     assert yaxis_model.name[0].name == "name"
     assert yaxis_model.name[0].value == "y axis"
-    assert yaxis_model.name[0].type == "str"
+    assert yaxis_model.name[0].type == "builtins.str"
 
 
 def test_matrix_coord() -> None:
     assert image_model.coord[0].name == "mask"
     assert image_model.coord[0].type == {"dims": ("x", "y"), "dtype": "bool"}
+    assert image_model.coord[0].factory is None
+
     assert image_model.coord[1].name == "x"
-    assert image_model.coord[1].type == "test_datamodel.XAxis"
+    assert image_model.coord[1].type == {"dims": ("x",), "dtype": "int"}
+    assert image_model.coord[1].factory is XAxis
+
     assert image_model.coord[2].name == "y"
-    assert image_model.coord[2].type == "test_datamodel.YAxis"
+    assert image_model.coord[2].type == {"dims": ("y",), "dtype": "int"}
+    assert image_model.coord[2].factory is YAxis
 
 
 def test_matrix_data() -> None:
@@ -104,8 +111,13 @@ def test_matrix_data() -> None:
 
 def test_image_data() -> None:
     assert color_model.data[0].name == "red"
-    assert color_model.data[0].type == "test_datamodel.Image"
+    assert color_model.data[0].type == {"dims": ("x", "y"), "dtype": "float"}
+    assert color_model.data[0].factory is Image
+
     assert color_model.data[1].name == "green"
-    assert color_model.data[1].type == "test_datamodel.Image"
+    assert color_model.data[1].type == {"dims": ("x", "y"), "dtype": "float"}
+    assert color_model.data[1].factory is Image
+
     assert color_model.data[2].name == "blue"
-    assert color_model.data[2].type == "test_datamodel.Image"
+    assert color_model.data[2].type == {"dims": ("x", "y"), "dtype": "float"}
+    assert color_model.data[2].factory is Image
diff --git a/tests/test_typing.py b/tests/test_typing.py
index 53f2e88..bcd9a23 100644
--- a/tests/test_typing.py
+++ b/tests/test_typing.py
@@ -8,7 +8,7 @@ from typing_extensions import Literal
 
 
 # submodules
-from xarray_dataclasses.typing import Data, get_dims, get_dtype
+from xarray_dataclasses.typing import Data, get_dims, get_dtype, unannotate
 
 
 # type hints
@@ -36,9 +36,9 @@ testdata_dtype = [
 # test functions
 @mark.parametrize("hint, dims", testdata_dims)
 def test_get_dims(hint: Any, dims: Any) -> None:
-    assert get_dims(hint) == dims
+    assert get_dims(unannotate(hint)) == dims
 
 
 @mark.parametrize("hint, dtype", testdata_dtype)
 def test_get_dtype(hint: Any, dtype: Any) -> None:
-    assert get_dtype(hint) == dtype
+    assert get_dtype(unannotate(hint)) == dtype