#163 Merge pull request from astropenguin/astropenguin/issue156

Add specs module

Author: astropenguin

tests/test_specs.py

@@ -0,0 +1,181 @@
+# standard library
+from dataclasses import MISSING, dataclass
+from typing import Tuple
+
+
+# dependencies
+import numpy as np
+import xarray as xr
+from typing_extensions import Annotated as Ann
+from typing_extensions import Literal as L
+from xarray_dataclasses.specs import DataOptions, DataSpec
+from xarray_dataclasses.typing import Attr, Coordof, Data, Name
+
+
+# type hints
+DataDims = Tuple[L["lon"], L["lat"], L["time"]]
+
+
+# test datasets
+@dataclass
+class Lon:
+    """Specification of relative longitude."""
+
+    data: Data[L["lon"], float]
+    units: Attr[str] = "deg"
+    name: Name[str] = "Relative longitude"
+
+
+@dataclass
+class Lat:
+    """Specification of relative latitude."""
+
+    data: Data[L["lat"], float]
+    units: Attr[str] = "m"
+    name: Name[str] = "Relative latitude"
+
+
+@dataclass
+class Time:
+    """Specification of time."""
+
+    data: Data[L["time"], L["datetime64[ns]"]]
+    name: Name[str] = "Time in UTC"
+
+
+@dataclass
+class Weather:
+    """Time-series spatial weather information at a location."""
+
+    temperature: Ann[Data[DataDims, float], "Temperature"]
+    humidity: Ann[Data[DataDims, float], "Humidity"]
+    wind_speed: Ann[Data[DataDims, float], "Wind speed"]
+    wind_direction: Ann[Data[DataDims, float], "Wind direction"]
+    lon: Coordof[Lon]
+    lat: Coordof[Lat]
+    time: Coordof[Time]
+    location: Attr[str] = "Tokyo"
+    longitude: Attr[float] = 139.69167
+    latitude: Attr[float] = 35.68944
+    name: Name[str] = "weather"
+
+
+# test functions
+def test_temperature() -> None:
+    spec = DataSpec.from_dataclass(Weather).specs.of_data["temperature"]
+
+    assert spec.name == "Temperature"
+    assert spec.role == "data"
+    assert spec.dims == ("lon", "lat", "time")
+    assert spec.dtype == np.dtype("f8")
+    assert spec.default is MISSING
+    assert spec.origin is None
+
+
+def test_humidity() -> None:
+    spec = DataSpec.from_dataclass(Weather).specs.of_data["humidity"]
+
+    assert spec.name == "Humidity"
+    assert spec.role == "data"
+    assert spec.dims == ("lon", "lat", "time")
+    assert spec.dtype == np.dtype("f8")
+    assert spec.default is MISSING
+    assert spec.origin is None
+
+
+def test_wind_speed() -> None:
+    spec = DataSpec.from_dataclass(Weather).specs.of_data["wind_speed"]
+
+    assert spec.name == "Wind speed"
+    assert spec.role == "data"
+    assert spec.dims == ("lon", "lat", "time")
+    assert spec.dtype == np.dtype("f8")
+    assert spec.default is MISSING
+    assert spec.origin is None
+
+
+def test_wind_direction() -> None:
+    spec = DataSpec.from_dataclass(Weather).specs.of_data["wind_direction"]
+
+    assert spec.name == "Wind direction"
+    assert spec.role == "data"
+    assert spec.dims == ("lon", "lat", "time")
+    assert spec.dtype == np.dtype("f8")
+    assert spec.default is MISSING
+    assert spec.origin is None
+
+
+def test_lon() -> None:
+    spec = DataSpec.from_dataclass(Weather).specs.of_coord["lon"]
+
+    assert spec.name == "Relative longitude"
+    assert spec.role == "coord"
+    assert spec.dims == ("lon",)
+    assert spec.dtype == np.dtype("f8")
+    assert spec.default is MISSING
+    assert spec.origin is Lon
+
+
+def test_lat() -> None:
+    spec = DataSpec.from_dataclass(Weather).specs.of_coord["lat"]
+
+    assert spec.name == "Relative latitude"
+    assert spec.role == "coord"
+    assert spec.dims == ("lat",)
+    assert spec.dtype == np.dtype("f8")
+    assert spec.default is MISSING
+    assert spec.origin is Lat
+
+
+def test_time() -> None:
+    spec = DataSpec.from_dataclass(Weather).specs.of_coord["time"]
+
+    assert spec.name == "Time in UTC"
+    assert spec.role == "coord"
+    assert spec.dims == ("time",)
+    assert spec.dtype == np.dtype("M8[ns]")
+    assert spec.default is MISSING
+    assert spec.origin is Time
+
+
+def test_location() -> None:
+    spec = DataSpec.from_dataclass(Weather).specs.of_attr["location"]
+
+    assert spec.name == "location"
+    assert spec.role == "attr"
+    assert spec.type is str
+    assert spec.default == "Tokyo"
+
+
+def test_longitude() -> None:
+    spec = DataSpec.from_dataclass(Weather).specs.of_attr["longitude"]
+
+    assert spec.name == "longitude"
+    assert spec.role == "attr"
+    assert spec.type is float
+    assert spec.default == 139.69167
+
+
+def test_latitude() -> None:
+    spec = DataSpec.from_dataclass(Weather).specs.of_attr["latitude"]
+
+    assert spec.name == "latitude"
+    assert spec.role == "attr"
+    assert spec.type is float
+    assert spec.default == 35.68944
+
+
+def test_name() -> None:
+    spec = DataSpec.from_dataclass(Weather).specs.of_name["name"]
+
+    assert spec.name == "name"
+    assert spec.role == "name"
+    assert spec.type is str
+    assert spec.default == "weather"
+
+
+def test_dataoptions() -> None:
+    options = DataOptions(xr.DataArray)
+
+    assert DataSpec().options.factory is type(None)
+    assert DataSpec(options=options).options.factory is xr.DataArray

tests/test_typing.py

@@ -14,12 +14,12 @@
     Attr,
     Coord,
     Data,
-    FType,
     Name,
+    Role,
     get_dims,
     get_dtype,
-    get_ftype,
     get_name,
+    get_role,
 )
 
 
@@ -54,24 +54,6 @@
     (Union[Ann[Data[Any, float], "data"], Ann[Any, "any"]], np.dtype("f8")),
 ]
 
-testdata_ftype = [
-    (Attr[Any], FType.ATTR),
-    (Data[Any, Any], FType.DATA),
-    (Coord[Any, Any], FType.COORD),
-    (Name[Any], FType.NAME),
-    (Any, FType.OTHER),
-    (Ann[Attr[Any], "attr"], FType.ATTR),
-    (Ann[Data[Any, Any], "data"], FType.DATA),
-    (Ann[Coord[Any, Any], "coord"], FType.COORD),
-    (Ann[Name[Any], "name"], FType.NAME),
-    (Ann[Any, "other"], FType.OTHER),
-    (Union[Ann[Attr[Any], "attr"], Ann[Any, "any"]], FType.ATTR),
-    (Union[Ann[Data[Any, Any], "data"], Ann[Any, "any"]], FType.DATA),
-    (Union[Ann[Coord[Any, Any], "coord"], Ann[Any, "any"]], FType.COORD),
-    (Union[Ann[Name[Any], "name"], Ann[Any, "any"]], FType.NAME),
-    (Union[Ann[Any, "other"], Ann[Any, "any"]], FType.OTHER),
-]
-
 testdata_name = [
     (Attr[Any], None),
     (Data[Any, Any], None),
@@ -90,6 +72,24 @@
     (Union[Ann[Any, "other"], Ann[Any, "any"]], None),
 ]
 
+testdata_role = [
+    (Attr[Any], Role.ATTR),
+    (Data[Any, Any], Role.DATA),
+    (Coord[Any, Any], Role.COORD),
+    (Name[Any], Role.NAME),
+    (Any, Role.OTHER),
+    (Ann[Attr[Any], "attr"], Role.ATTR),
+    (Ann[Data[Any, Any], "data"], Role.DATA),
+    (Ann[Coord[Any, Any], "coord"], Role.COORD),
+    (Ann[Name[Any], "name"], Role.NAME),
+    (Ann[Any, "other"], Role.OTHER),
+    (Union[Ann[Attr[Any], "attr"], Ann[Any, "any"]], Role.ATTR),
+    (Union[Ann[Data[Any, Any], "data"], Ann[Any, "any"]], Role.DATA),
+    (Union[Ann[Coord[Any, Any], "coord"], Ann[Any, "any"]], Role.COORD),
+    (Union[Ann[Name[Any], "name"], Ann[Any, "any"]], Role.NAME),
+    (Union[Ann[Any, "other"], Ann[Any, "any"]], Role.OTHER),
+]
+
 
 # test functions
 @mark.parametrize("tp, dims", testdata_dims)
@@ -102,11 +102,11 @@ def test_get_dtype(tp: Any, dtype: Any) -> None:
     assert get_dtype(tp) == dtype
 
 
-@mark.parametrize("tp, ftype", testdata_ftype)
-def test_get_ftype(tp: Any, ftype: Any) -> None:
-    assert get_ftype(tp) == ftype
-
-
 @mark.parametrize("tp, name", testdata_name)
 def test_get_name(tp: Any, name: Any) -> None:
     assert get_name(tp) == name
+
+
+@mark.parametrize("tp, role", testdata_role)
+def test_get_role(tp: Any, role: Any) -> None:
+    assert get_role(tp) == role

xarray_dataclasses/__init__.py

@@ -7,6 +7,7 @@
 from .dataset import *
 from .datamodel import *
 from .dataoptions import *
+from .specs import *
 from .typing import *
 
 

xarray_dataclasses/datamodel.py

@@ -17,16 +17,16 @@
 from .typing import (
     AnyDType,
     AnyField,
-    DataClass,
     AnyXarray,
+    DataClass,
     Dims,
-    FType,
+    Role,
     get_annotated,
     get_dataclass,
     get_dims,
     get_dtype,
-    get_ftype,
     get_name,
+    get_role,
 )
 
 
@@ -209,29 +209,29 @@ def eval_dataclass(dataclass: AnyDataClass[PInit]) -> None:
 
 def get_entry(field: AnyField, value: Any) -> Optional[AnyEntry]:
     """Create an entry from a field and its value."""
-    ftype = get_ftype(field.type)
+    role = get_role(field.type)
     name = get_name(field.type, field.name)
 
-    if ftype is FType.ATTR or ftype is FType.NAME:
+    if role is Role.ATTR or role is Role.NAME:
         return AttrEntry(
             name=name,
-            tag=ftype.value,
+            tag=role.value,
             value=value,
             type=get_annotated(field.type),
         )
 
-    if ftype is FType.COORD or ftype is FType.DATA:
+    if role is Role.COORD or role is Role.DATA:
         try:
             return DataEntry(
                 name=name,
-                tag=ftype.value,
+                tag=role.value,
                 base=get_dataclass(field.type),
                 value=value,
             )
         except TypeError:
             return DataEntry(
                 name=name,
-                tag=ftype.value,
+                tag=role.value,
                 dims=get_dims(field.type),
                 dtype=get_dtype(field.type),
                 value=value,