Merge pull request #138 from astropenguin/astropenguin/issue137

Update data model

Author: astropenguin

poetry.lock

@@ -12,17 +12,18 @@ documentation = "https://astropenguin.github.io/xarray-dataclasses/"
 [tool.poetry.dependencies]
 python = ">=3.7.1, <3.10"
 morecopy = "^0.2"
+more-itertools = "^8.12"
 numpy = "^1.19"
 typing-extensions = "^3.10"
 xarray = ">=0.18, <0.30"
 
 [tool.poetry.dev-dependencies]
-black = "^21.12b"
+black = "^22.1"
 ipython = "^7.31"
 myst-parser = "^0.16"
-pydata-sphinx-theme = "^0.7"
+pydata-sphinx-theme = "^0.8"
 pytest = "^6.2"
-sphinx = "^4.3"
+sphinx = "^4.4"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

pyproject.toml

@@ -53,73 +53,100 @@ class ColorImage:
 
 # test functions
 def test_xaxis_attr() -> None:
-    item = next(iter(xaxis_model.attr.values()))
-    assert item.name == "units"
-    assert item.value == "pixel"
-    assert item.type == "builtins.str"
+    units = xaxis_model.attrs[0]
+    assert units.name == "units"
+    assert units.tag == "attr"
+    assert units.type is str
+    assert units.value == "pixel"
+    assert units.cast == False
 
 
 def test_xaxis_data() -> None:
-    item = next(iter(xaxis_model.data.values()))
-    assert item.name == "data"
-    assert item.type == {"dims": ("x",), "dtype": "int"}
-    assert item.factory is None
+    data = xaxis_model.data_vars[0]
+    assert data.name == "data"
+    assert data.tag == "data"
+    assert data.dims == ("x",)
+    assert data.dtype == "int"
+    assert data.base is None
+    assert data.cast == True
 
 
 def test_yaxis_attr() -> None:
-    item = next(iter(yaxis_model.attr.values()))
-    assert item.name == "units"
-    assert item.value == "pixel"
-    assert item.type == "builtins.str"
+    units = yaxis_model.attrs[0]
+    assert units.name == "units"
+    assert units.tag == "attr"
+    assert units.type is str
+    assert units.value == "pixel"
+    assert units.cast == False
 
 
 def test_yaxis_data() -> None:
-    item = next(iter(yaxis_model.data.values()))
-    assert item.name == "data"
-    assert item.type == {"dims": ("y",), "dtype": "int"}
-    assert item.factory is None
+    data = yaxis_model.data_vars[0]
+    assert data.name == "data"
+    assert data.tag == "data"
+    assert data.dims == ("y",)
+    assert data.dtype == "int"
+    assert data.base is None
+    assert data.cast == True
 
 
 def test_image_coord() -> None:
-    items = iter(image_model.coord.values())
-
-    item = next(items)
-    assert item.name == "mask"
-    assert item.type == {"dims": ("x", "y"), "dtype": "bool"}
-    assert item.factory is None
-
-    item = next(items)
-    assert item.name == "x"
-    assert item.type == {"dims": ("x",), "dtype": "int"}
-    assert item.factory is XAxis
-
-    item = next(items)
-    assert item.name == "y"
-    assert item.type == {"dims": ("y",), "dtype": "int"}
-    assert item.factory is YAxis
+    mask = image_model.coords[0]
+    assert mask.name == "mask"
+    assert mask.tag == "coord"
+    assert mask.dims == ("x", "y")
+    assert mask.dtype == "bool"
+    assert mask.base is None
+    assert mask.cast == True
+
+    x = image_model.coords[1]
+    assert x.name == "x"
+    assert x.tag == "coord"
+    assert x.dims == ("x",)
+    assert x.dtype == "int"
+    assert x.base is XAxis
+    assert x.cast == True
+
+    y = image_model.coords[2]
+    assert y.name == "y"
+    assert y.tag == "coord"
+    assert y.dims == ("y",)
+    assert y.dtype == "int"
+    assert y.base is YAxis
+    assert y.cast == True
 
 
 def test_image_data() -> None:
-    item = next(iter(image_model.data.values()))
-    assert item.name == "data"
-    assert item.type == {"dims": ("x", "y"), "dtype": "float"}
-    assert item.factory is None
+    data = image_model.data_vars[0]
+    assert data.name == "data"
+    assert data.tag == "data"
+    assert data.dims == ("x", "y")
+    assert data.dtype == "float"
+    assert data.base is None
+    assert data.cast == True
 
 
 def test_color_data() -> None:
-    items = iter(color_model.data.values())
-
-    item = next(items)
-    assert item.name == "red"
-    assert item.type == {"dims": ("x", "y"), "dtype": "float"}
-    assert item.factory is Image
-
-    item = next(items)
-    assert item.name == "green"
-    assert item.type == {"dims": ("x", "y"), "dtype": "float"}
-    assert item.factory is Image
-
-    item = next(items)
-    assert item.name == "blue"
-    assert item.type == {"dims": ("x", "y"), "dtype": "float"}
-    assert item.factory is Image
+    red = color_model.data_vars[0]
+    assert red.name == "red"
+    assert red.tag == "data"
+    assert red.dims == ("x", "y")
+    assert red.dtype == "float"
+    assert red.base is Image
+    assert red.cast == True
+
+    green = color_model.data_vars[1]
+    assert green.name == "green"
+    assert green.tag == "data"
+    assert green.dims == ("x", "y")
+    assert green.dtype == "float"
+    assert green.base is Image
+    assert green.cast == True
+
+    blue = color_model.data_vars[2]
+    assert blue.name == "blue"
+    assert blue.tag == "data"
+    assert blue.dims == ("x", "y")
+    assert blue.dtype == "float"
+    assert blue.base is Image
+    assert blue.cast == True

tests/test_datamodel.py

@@ -1,44 +1,87 @@
 # standard library
-from typing import Any, Tuple
+from typing import Any, Optional, Tuple, Union
 
 
 # third-party packages
 from pytest import mark
-from typing_extensions import Literal
+from typing_extensions import Annotated, Literal
 
 
 # submodules
-from xarray_dataclasses.typing import Data, get_dims, get_dtype, unannotate
+from xarray_dataclasses.typing import (
+    ArrayLike,
+    Attr,
+    Coord,
+    Data,
+    Name,
+    get_dims,
+    get_dtype,
+    get_field_type,
+    get_repr_type,
+)
 
 
 # type hints
-Int = Literal["int"]
+Int64 = Literal["int64"]
+NoneType = type(None)
 X = Literal["x"]
 Y = Literal["y"]
 
 
 # test datasets
 testdata_dims = [
-    (Data[X, Any], ("x",)),
-    (Data[Tuple[()], Any], ()),
-    (Data[Tuple[X], Any], ("x",)),
-    (Data[Tuple[X, Y], Any], ("x", "y")),
+    (X, ("x",)),
+    (Tuple[()], ()),
+    (Tuple[X], ("x",)),
+    (Tuple[X, Y], ("x", "y")),
+    (ArrayLike[X, Any], ("x",)),
+    (ArrayLike[Tuple[()], Any], ()),
+    (ArrayLike[Tuple[X], Any], ("x",)),
+    (ArrayLike[Tuple[X, Y], Any], ("x", "y")),
 ]
 
 testdata_dtype = [
-    (Data[X, Any], None),
-    (Data[X, None], None),
-    (Data[X, int], "int"),
-    (Data[X, Int], "int"),
+    (Any, None),
+    (NoneType, None),
+    (Int64, "int64"),
+    (int, "int"),
+    (ArrayLike[Any, Any], None),
+    (ArrayLike[Any, NoneType], None),
+    (ArrayLike[Any, Int64], "int64"),
+    (ArrayLike[Any, int], "int"),
+]
+
+testdata_field_type = [
+    (Attr[Any], "attr"),
+    (Coord[Any, Any], "coord"),
+    (Data[Any, Any], "data"),
+    (Name[Any], "name"),
+]
+
+testdata_repr_type = [
+    (int, int),
+    (Annotated[int, "annotation"], int),
+    (Union[int, float], int),
+    (Optional[int], int),
 ]
 
 
 # test functions
-@mark.parametrize("hint, dims", testdata_dims)
-def test_get_dims(hint: Any, dims: Any) -> None:
-    assert get_dims(unannotate(hint)) == dims
+@mark.parametrize("type_, dims", testdata_dims)
+def test_get_dims(type_: Any, dims: Any) -> None:
+    assert get_dims(type_) == dims
+
+
+@mark.parametrize("type_, dtype", testdata_dtype)
+def test_get_dtype(type_: Any, dtype: Any) -> None:
+    assert get_dtype(type_) == dtype
+
+
+@mark.parametrize("type_, field_type", testdata_field_type)
+def test_get_field_type(type_: Any, field_type: Any) -> None:
+    assert get_field_type(type_).value == field_type
 
 
-@mark.parametrize("hint, dtype", testdata_dtype)
-def test_get_dtype(hint: Any, dtype: Any) -> None:
-    assert get_dtype(unannotate(hint)) == dtype
+@mark.parametrize("type_, repr_type", testdata_repr_type)
+def test_get_repr_type(type_: Any, repr_type: Any) -> None:
+    assert get_repr_type(type_) == repr_type

tests/test_typing.py

@@ -104,23 +104,21 @@ def asdataarray(
             dataoptions = DataOptions(xr.DataArray)
 
     model = DataModel.from_dataclass(dataclass)
-    item = next(iter(model.data.values()))
-    dataarray = dataoptions.factory(item(reference))
+    dataarray = dataoptions.factory(model.data_vars[0](reference))
 
-    for item in model.coord.values():
-        if item.name in dataarray.dims:
-            dataarray.coords[item.name] = item(dataarray)
+    for entry in model.coords:
+        if entry.name in dataarray.dims:
+            dataarray.coords[entry.name] = entry(dataarray)
 
-    for item in model.coord.values():
-        if item.name not in dataarray.dims:
-            dataarray.coords[item.name] = item(dataarray)
+    for entry in model.coords:
+        if entry.name not in dataarray.dims:
+            dataarray.coords[entry.name] = entry(dataarray)
 
-    for item in model.attr.values():
-        dataarray.attrs[item.name] = item()
+    for entry in model.attrs:
+        dataarray.attrs[entry.name] = entry()
 
-    if model.name:
-        item = next(iter(model.name.values()))
-        dataarray.name = item()
+    if model.names:
+        dataarray.name = model.names[0]()
 
     return dataarray
 
@@ -213,12 +211,12 @@ def shaped(
 
         """
         model = DataModel.from_dataclass(cls)
-        name, item = next(iter(model.data.items()))
+        key, entry = model.data_vars_items[0]
 
         if isinstance(shape, dict):
-            shape = tuple(shape[dim] for dim in item.type["dims"])
+            shape = tuple(shape[dim] for dim in entry.dims)
 
-        return asdataarray(cls(**{name: func(shape)}, **kwargs))
+        return asdataarray(cls(**{key: func(shape)}, **kwargs))
 
     @overload
     @classmethod