#91 Use dict in DataModel

Author: astropenguin

tests/test_datamodel.py

@@ -53,57 +53,73 @@ class ColorImage:
 
 # test functions
 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 == "builtins.str"
+    item = next(iter(xaxis_model.attr.values()))
+    assert item.key == "units"
+    assert item.value == "pixel"
+    assert item.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
+    item = next(iter(xaxis_model.data.values()))
+    assert item.key == "data"
+    assert item.type == {"dims": ("x",), "dtype": "int"}
+    assert item.factory is None
 
 
 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 == "builtins.str"
+    item = next(iter(yaxis_model.attr.values()))
+    assert item.key == "units"
+    assert item.value == "pixel"
+    assert item.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
+    item = next(iter(yaxis_model.data.values()))
+    assert item.key == "data"
+    assert item.type == {"dims": ("y",), "dtype": "int"}
+    assert item.factory is None
 
 
-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
+def test_image_coord() -> None:
+    items = iter(image_model.coord.values())
 
-    assert image_model.coord[1].name == "x"
-    assert image_model.coord[1].type == {"dims": ("x",), "dtype": "int"}
-    assert image_model.coord[1].factory is XAxis
+    item = next(items)
+    assert item.key == "mask"
+    assert item.type == {"dims": ("x", "y"), "dtype": "bool"}
+    assert item.factory is None
 
-    assert image_model.coord[2].name == "y"
-    assert image_model.coord[2].type == {"dims": ("y",), "dtype": "int"}
-    assert image_model.coord[2].factory is YAxis
+    item = next(items)
+    assert item.key == "x"
+    assert item.type == {"dims": ("x",), "dtype": "int"}
+    assert item.factory is XAxis
 
-
-def test_matrix_data() -> None:
-    assert image_model.data[0].name == "data"
-    assert image_model.data[0].type == {"dims": ("x", "y"), "dtype": "float"}
+    item = next(items)
+    assert item.key == "y"
+    assert item.type == {"dims": ("y",), "dtype": "int"}
+    assert item.factory is YAxis
 
 
 def test_image_data() -> None:
-    assert color_model.data[0].name == "red"
-    assert color_model.data[0].type == {"dims": ("x", "y"), "dtype": "float"}
-    assert color_model.data[0].factory is Image
+    item = next(iter(image_model.data.values()))
+    assert item.key == "data"
+    assert item.type == {"dims": ("x", "y"), "dtype": "float"}
+    assert item.factory is None
+
+
+def test_color_data() -> None:
+    items = iter(color_model.data.values())
+
+    item = next(items)
+    assert item.key == "red"
+    assert item.type == {"dims": ("x", "y"), "dtype": "float"}
+    assert item.factory is Image
 
-    assert color_model.data[1].name == "green"
-    assert color_model.data[1].type == {"dims": ("x", "y"), "dtype": "float"}
-    assert color_model.data[1].factory is Image
+    item = next(items)
+    assert item.key == "green"
+    assert item.type == {"dims": ("x", "y"), "dtype": "float"}
+    assert item.factory is Image
 
-    assert color_model.data[2].name == "blue"
-    assert color_model.data[2].type == {"dims": ("x", "y"), "dtype": "float"}
-    assert color_model.data[2].factory is Image
+    item = next(items)
+    assert item.key == "blue"
+    assert item.type == {"dims": ("x", "y"), "dtype": "float"}
+    assert item.factory is Image

xarray_dataclasses/dataarray.py

@@ -113,21 +113,23 @@ def asdataarray(
         pass
 
     model = DataModel.from_dataclass(dataclass)
-    dataarray = dataoptions.factory(model.data[0](reference))
+    item = next(iter(model.data.values()))
+    dataarray = dataoptions.factory(item(reference))
 
-    for coord in model.coord:
-        if coord.name in dataarray.dims:
-            dataarray.coords.update({coord.name: coord(dataarray)})
+    for item in model.coord.values():
+        if item.key in dataarray.dims:
+            dataarray.coords[item.key] = item(dataarray)
 
-    for coord in model.coord:
-        if coord.name not in dataarray.dims:
-            dataarray.coords.update({coord.name: coord(dataarray)})
+    for item in model.coord.values():
+        if item.key not in dataarray.dims:
+            dataarray.coords[item.key] = item(dataarray)
 
-    for attr in model.attr:
-        dataarray.attrs.update({attr.name: attr()})
+    for item in model.attr.values():
+        dataarray.attrs[item.key] = item()
 
-    for name in model.name:
-        dataarray.name = name()
+    if model.name:
+        item = next(iter(model.name.values()))
+        dataarray.name = item()
 
     return dataarray
 
@@ -175,11 +177,10 @@ def empty(
 
         """
         model = DataModel.from_dataclass(cls)
-        name = model.data[0].name
-        dims = model.data[0].type["dims"]
+        name, item = next(iter(model.data.items()))
 
         if isinstance(shape, dict):
-            shape = tuple(shape[dim] for dim in dims)
+            shape = tuple(shape[dim] for dim in item.type["dims"])
 
         data = np.empty(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
@@ -204,11 +205,10 @@ def zeros(
 
         """
         model = DataModel.from_dataclass(cls)
-        name = model.data[0].name
-        dims = model.data[0].type["dims"]
+        name, item = next(iter(model.data.items()))
 
         if isinstance(shape, dict):
-            shape = tuple(shape[dim] for dim in dims)
+            shape = tuple(shape[dim] for dim in item.type["dims"])
 
         data = np.zeros(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
@@ -233,11 +233,10 @@ def ones(
 
         """
         model = DataModel.from_dataclass(cls)
-        name = model.data[0].name
-        dims = model.data[0].type["dims"]
+        name, item = next(iter(model.data.items()))
 
         if isinstance(shape, dict):
-            shape = tuple(shape[dim] for dim in dims)
+            shape = tuple(shape[dim] for dim in item.type["dims"])
 
         data = np.ones(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
@@ -264,11 +263,10 @@ def full(
 
         """
         model = DataModel.from_dataclass(cls)
-        name = model.data[0].name
-        dims = model.data[0].type["dims"]
+        name, item = next(iter(model.data.items()))
 
         if isinstance(shape, dict):
-            shape = tuple(shape[dim] for dim in dims)
+            shape = tuple(shape[dim] for dim in item.type["dims"])
 
         data = np.full(shape, fill_value, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))

xarray_dataclasses/datamodel.py

@@ -3,7 +3,7 @@
 
 # standard library
 from dataclasses import Field, dataclass, field, is_dataclass
-from typing import Any, Hashable, List, Optional, Type, cast
+from typing import Any, Dict, Hashable, Optional, Type, cast
 
 
 # dependencies
@@ -37,7 +37,7 @@
 class Data:
     """Field model for data-related fields."""
 
-    name: Hashable
+    key: Hashable
     """Name of the field."""
 
     value: Any
@@ -77,18 +77,20 @@ def from_field(cls, field: Field[Any], value: Any, of: bool) -> "Data":
 
         dataclass = get_inner(hint, 0)
         model = DataModel.from_dataclass(dataclass)
+        data_item = next(iter(model.data.values()))
 
         if not model.name:
-            return cls(field.name, value, model.data[0].type, dataclass)
+            return cls(field.name, value, data_item.type, dataclass)
         else:
-            return cls(model.name[0].value, value, model.data[0].type, dataclass)
+            name_item = next(iter(model.name.values()))
+            return cls(name_item.value, value, data_item.type, dataclass)
 
 
 @dataclass(frozen=True)
 class General:
     """Field model for general fields."""
 
-    name: Hashable
+    key: Hashable
     """Name of the field."""
 
     value: Any
@@ -123,16 +125,16 @@ def from_field(cls, field: Field[Any], value: Any) -> "General":
 class DataModel:
     """Model for dataclasses or their objects."""
 
-    attr: List[General] = field(default_factory=list)
+    attr: Dict[str, General] = field(default_factory=dict)
     """Model of the attribute fields."""
 
-    coord: List[Data] = field(default_factory=list)
+    coord: Dict[str, Data] = field(default_factory=dict)
     """Model of the coordinate fields."""
 
-    data: List[Data] = field(default_factory=list)
+    data: Dict[str, Data] = field(default_factory=dict)
     """Model of the data fields."""
 
-    name: List[General] = field(default_factory=list)
+    name: Dict[str, General] = field(default_factory=dict)
     """Model of the name fields."""
 
     @classmethod
@@ -141,21 +143,21 @@ def from_dataclass(cls, dataclass: DataClass) -> "DataModel":
         model = cls()
         eval_field_types(dataclass)
 
-        for field_ in dataclass.__dataclass_fields__.values():
-            value = getattr(dataclass, field_.name, field_.default)
-
-            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, False))
-            elif FieldType.COORDOF.annotates(field_.type):
-                model.coord.append(Data.from_field(field_, value, True))
-            elif FieldType.DATA.annotates(field_.type):
-                model.data.append(Data.from_field(field_, value, False))
-            elif FieldType.DATAOF.annotates(field_.type):
-                model.data.append(Data.from_field(field_, value, True))
-            elif FieldType.NAME.annotates(field_.type):
-                model.name.append(General.from_field(field_, value))
+        for field in dataclass.__dataclass_fields__.values():
+            value = getattr(dataclass, field.name, field.default)
+
+            if FieldType.ATTR.annotates(field.type):
+                model.attr[field.name] = General.from_field(field, value)
+            elif FieldType.COORD.annotates(field.type):
+                model.coord[field.name] = Data.from_field(field, value, False)
+            elif FieldType.COORDOF.annotates(field.type):
+                model.coord[field.name] = Data.from_field(field, value, True)
+            elif FieldType.DATA.annotates(field.type):
+                model.data[field.name] = Data.from_field(field, value, False)
+            elif FieldType.DATAOF.annotates(field.type):
+                model.data[field.name] = Data.from_field(field, value, True)
+            elif FieldType.NAME.annotates(field.type):
+                model.name[field.name] = General.from_field(field, value)
 
         return model
 
@@ -165,9 +167,9 @@ def eval_field_types(dataclass: DataClass) -> None:
     """Evaluate field types of a dataclass or its object."""
     hints = get_type_hints(dataclass, include_extras=True)  # type: ignore
 
-    for field_ in dataclass.__dataclass_fields__.values():
-        if isinstance(field_.type, str):
-            field_.type = hints[field_.name]
+    for field in dataclass.__dataclass_fields__.values():
+        if isinstance(field.type, str):
+            field.type = hints[field.name]
 
 
 def typedarray(

xarray_dataclasses/dataset.py

@@ -115,19 +115,19 @@ def asdataset(
     model = DataModel.from_dataclass(dataclass)
     dataset = dataoptions.factory()
 
-    for data in model.data:
-        dataset.update({data.name: data(reference)})
+    for item in model.data.values():
+        dataset[item.key] = item(reference)
 
-    for coord in model.coord:
-        if coord.name in dataset.dims:
-            dataset.coords.update({coord.name: coord(dataset)})
+    for item in model.coord.values():
+        if item.key in dataset.dims:
+            dataset.coords[item.key] = item(dataset)
 
-    for coord in model.coord:
-        if coord.name not in dataset.dims:
-            dataset.coords.update({coord.name: coord(dataset)})
+    for item in model.coord.values():
+        if item.key not in dataset.dims:
+            dataset.coords[item.key] = item(dataset)
 
-    for attr in model.attr:
-        dataset.attrs.update({attr.name: attr()})
+    for item in model.attr.values():
+        dataset.attrs[item.key] = item()
 
     return dataset
 
@@ -177,9 +177,9 @@ def empty(
         model = DataModel.from_dataclass(cls)
         data_vars: Dict[str, Any] = {}
 
-        for data in model.data:
-            shape = tuple(sizes[dim] for dim in data.type["dims"])
-            data_vars[data.name] = np.empty(shape, order=order)
+        for name, item in model.data.items():
+            shape = tuple(sizes[dim] for dim in item.type["dims"])
+            data_vars[name] = np.empty(shape, order=order)
 
         return asdataset(cls(**data_vars, **kwargs))
 
@@ -205,9 +205,9 @@ def zeros(
         model = DataModel.from_dataclass(cls)
         data_vars: Dict[str, Any] = {}
 
-        for data in model.data:
-            shape = tuple(sizes[dim] for dim in data.type["dims"])
-            data_vars[data.name] = np.zeros(shape, order=order)
+        for name, item in model.data.items():
+            shape = tuple(sizes[dim] for dim in item.type["dims"])
+            data_vars[name] = np.zeros(shape, order=order)
 
         return asdataset(cls(**data_vars, **kwargs))
 
@@ -233,9 +233,9 @@ def ones(
         model = DataModel.from_dataclass(cls)
         data_vars: Dict[str, Any] = {}
 
-        for data in model.data:
-            shape = tuple(sizes[dim] for dim in data.type["dims"])
-            data_vars[data.name] = np.ones(shape, order=order)
+        for name, item in model.data.items():
+            shape = tuple(sizes[dim] for dim in item.type["dims"])
+            data_vars[name] = np.ones(shape, order=order)
 
         return asdataset(cls(**data_vars, **kwargs))
 
@@ -263,8 +263,8 @@ def full(
         model = DataModel.from_dataclass(cls)
         data_vars: Dict[str, Any] = {}
 
-        for data in model.data:
-            shape = tuple(sizes[dim] for dim in data.type["dims"])
-            data_vars[data.name] = np.full(shape, fill_value, order=order)
+        for name, item in model.data.items():
+            shape = tuple(sizes[dim] for dim in item.type["dims"])
+            data_vars[name] = np.full(shape, fill_value, order=order)
 
         return asdataset(cls(**data_vars, **kwargs))