Merge pull request #118 from astropenguin/astropenguin/issue91

Use name-field value as data-variable name

Author: astropenguin

tests/test_datamodel.py

@@ -9,7 +9,7 @@
 
 # submodules
 from xarray_dataclasses.datamodel import DataModel
-from xarray_dataclasses.typing import Attr, Coord, Coordof, Data, Dataof, Name
+from xarray_dataclasses.typing import Attr, Coord, Coordof, Data, Dataof
 
 
 # type hints
@@ -22,14 +22,12 @@
 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
@@ -55,69 +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.name == "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
-
-
-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 == "builtins.str"
+    item = next(iter(xaxis_model.data.values()))
+    assert item.name == "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.name == "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.name == "data"
+    assert item.type == {"dims": ("y",), "dtype": "int"}
+    assert item.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 == "builtins.str"
+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
 
-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
+    item = next(items)
+    assert item.name == "x"
+    assert item.type == {"dims": ("x",), "dtype": "int"}
+    assert item.factory is XAxis
 
-    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.name == "y"
+    assert item.type == {"dims": ("y",), "dtype": "int"}
+    assert item.factory is YAxis
 
-    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
 
+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
 
-def test_matrix_data() -> None:
-    assert image_model.data[0].name == "data"
-    assert image_model.data[0].type == {"dims": ("x", "y"), "dtype": "float"}
 
+def test_color_data() -> None:
+    items = iter(color_model.data.values())
 
-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(items)
+    assert item.name == "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.name == "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.name == "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.name in dataarray.dims:
+            dataarray.coords[item.name] = 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.name not in dataarray.dims:
+            dataarray.coords[item.name] = item(dataarray)
 
-    for attr in model.attr:
-        dataarray.attrs.update({attr.name: attr()})
+    for item in model.attr.values():
+        dataarray.attrs[item.name] = 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, 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: str
+    name: Hashable
     """Name of the field."""
 
     value: Any
@@ -71,23 +71,26 @@ def from_field(cls, field: Field[Any], value: Any, of: bool) -> "Data":
         """Create a field model from a dataclass field and a value."""
         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)
+        if not of:
+            type: DimsDtype = {"dims": get_dims(hint), "dtype": get_dtype(hint)}
+            return cls(field.name, value, type)
+
+        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, data_item.type, dataclass)
         else:
-            return cls(
-                field.name,
-                value,
-                {"dims": get_dims(hint), "dtype": get_dtype(hint)},
-            )
+            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: str
+    name: Hashable
     """Name of the field."""
 
     value: Any
@@ -122,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
@@ -140,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
 
@@ -164,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(