first pass at attrs and coords schemas

Author: Joseph Hamman

README.md

@@ -38,7 +38,7 @@ The basic usage is as follows:
 ```python
 import numpy as np
 import xarray as xr
-from xarray_schema import DataArraySchema, DatasetSchema
+from xarray_schema import DataArraySchema, DatasetSchema, CoordsSchema
 
 da = xr.DataArray(np.ones(4, dtype='i4'), dims=['x'], name='foo')
 
@@ -64,7 +64,9 @@ from xarray_schema.components import (
     ShapeSchema,
     NameSchema,
     ChunksSchema,
-    ArrayTypeSchema
+    ArrayTypeSchema,
+    AttrSchema,
+    AttrsSchema
 )
 
 # example constructions

tests/test_core.py

@@ -6,6 +6,8 @@
 from xarray_schema.base import SchemaError
 from xarray_schema.components import (
     ArrayTypeSchema,
+    AttrSchema,
+    AttrsSchema,
     ChunksSchema,
     DimsSchema,
     DTypeSchema,
@@ -52,6 +54,18 @@ def ds():
             [(((2, 2), (10,)), ('x', 'y'), (4, 10))],
             {'x': 2, 'y': -1},
         ),
+        (
+            AttrsSchema,
+            {'foo': AttrSchema(value='bar')},
+            [{'foo': 'bar'}],
+            {'foo': {'type': None, 'value': 'bar'}},
+        ),
+        (
+            AttrsSchema,
+            {'foo': AttrSchema(value=1)},
+            [{'foo': 1}],
+            {'foo': {'type': None, 'value': 1}},
+        ),
     ],
 )
 def test_component_schema(component, schema_args, validate, json):
@@ -65,6 +79,21 @@ def test_component_schema(component, schema_args, validate, json):
     assert isinstance(schema.to_json(), str)
 
 
+@pytest.mark.parametrize(
+    'type, value, validate, json',
+    [
+        (str, None, 'foo', {'type': str, 'value': None}),
+        (None, 'foo', 'foo', {'type': None, 'value': 'foo'}),
+        (str, 'foo', 'foo', {'type': str, 'value': 'foo'}),
+    ],
+)
+def test_attr_schema(type, value, validate, json):
+    schema = AttrSchema(type=type, value=value)
+    schema.validate(validate)
+    assert schema.json == json
+    # assert isinstance(schema.to_json(), str)
+
+
 @pytest.mark.parametrize(
     'component, schema_args, validate, match',
     [

xarray_schema/__init__.py

@@ -8,7 +8,7 @@
     NameSchema,
     ShapeSchema,
 )
-from .dataarray import DataArraySchema  # noqa: F401
+from .dataarray import CoordsSchema, DataArraySchema  # noqa: F401
 from .dataset import DatasetSchema  # noqa: F401
 
 try:

xarray_schema/components.py

@@ -1,7 +1,8 @@
 from collections.abc import Iterable
-from typing import Any, Dict, Hashable, Optional, Tuple, Union
+from typing import Any, Dict, Hashable, Mapping, Optional, Tuple, Union
 
 import numpy as np
+import numpy.typing as npt
 
 from .base import BaseSchema, SchemaError
 from .types import ChunksT, DimsT, ShapeT
@@ -11,13 +12,13 @@ class DTypeSchema(BaseSchema):
 
     _json_schema = {'type': 'string'}
 
-    def __init__(self, dtype: np.typing.DTypeLike) -> None:
+    def __init__(self, dtype: npt.DTypeLike) -> None:
         if dtype in [np.floating, np.integer, np.signedinteger, np.unsignedinteger, np.generic]:
             self.dtype = dtype
         else:
             self.dtype = np.dtype(dtype)
 
-    def validate(self, dtype: np.typing.DTypeLike) -> None:
+    def validate(self, dtype: npt.DTypeLike) -> None:
         '''Validate dtype
 
         Parameters
@@ -201,3 +202,67 @@ def validate(self, array: Any) -> None:
     @property
     def json(self) -> str:
         return str(self.array_type)
+
+
+class AttrSchema(BaseSchema):
+
+    _json_schema = {'type': 'object'}  # TODO: add type/value here
+
+    def __init__(self, type: Any = None, value: Any = None):
+        self.type = type
+        self.value = value
+
+    def validate(self, attr: Any):
+
+        if self.type is not None:
+            if not isinstance(attr, self.type):
+                SchemaError(f'attrs {attr} is not of type {self.type}')
+
+        if self.value is not None:
+            if self.value is not None and self.value != attr:
+                raise SchemaError(f'name {attr} != {self.value}')
+
+    @property
+    def json(self) -> str:
+        return {'type': self.type, 'value': self.value}
+
+
+class AttrsSchema(BaseSchema):
+
+    _json_schema = {'type': 'string'}
+
+    def __init__(
+        self, attrs: Mapping, require_all_keys: bool = True, allow_extra_keys: bool = True
+    ) -> None:
+        self.attrs = attrs
+        self.require_all_keys = require_all_keys
+        self.allow_extra_keys = allow_extra_keys
+
+    def validate(self, attrs: Any) -> None:
+        '''Validate attrs
+
+        Parameters
+        ----------
+        attrs : dict_like
+            attrs of the DataArray. `None` may be used as a wildcard value.
+        '''
+
+        if self.require_all_keys:
+            missing_keys = set(self.attrs) - set(attrs)
+            if missing_keys:
+                raise SchemaError(f'attrs has missing keys: {missing_keys}')
+
+        if not self.allow_extra_keys:
+            extra_keys = set(attrs) - set(self.attrs)
+            if extra_keys:
+                raise SchemaError(f'attrs has extra keys: {extra_keys}')
+
+        for key, attr_schema in self.attrs.items():
+            if key not in attrs:
+                raise SchemaError(f'key {key} not in attrs')
+            else:
+                attr_schema.validate(attrs[key])
+
+    @property
+    def json(self) -> dict:
+        return {k: v.json for k, v in self.attrs.items()}

xarray_schema/dataarray.py

@@ -1,11 +1,12 @@
-from typing import Any, Callable, Dict, List, Union
+from typing import Any, Callable, Dict, List, Mapping, Union
 
 import numpy as np
 import xarray as xr
 
-from .base import BaseSchema
+from .base import BaseSchema, SchemaError
 from .components import (
     ArrayTypeSchema,
+    AttrsSchema,
     ChunksSchema,
     DimsSchema,
     DTypeSchema,
@@ -43,7 +44,9 @@ class DataArraySchema(BaseSchema):
     _shape: Union[ShapeSchema, None]
     _dims: Union[DimsSchema, None]
     _name: Union[NameSchema, None]
+    _coords: Union[Any, None]
     _chunks: Union[ChunksSchema, None]
+    _attrs: Union[AttrsSchema, None]
     _array_type: Union[ArrayTypeSchema, None]
 
     def __init__(
@@ -55,7 +58,7 @@ def __init__(
         coords: Dict[str, Any] = None,
         chunks: Union[ChunksT, ChunksSchema] = None,
         array_type: Any = None,
-        attrs: Dict[str, Any] = None,
+        attrs: Mapping[str, Any] = None,
         checks: List[Callable] = None,
     ) -> None:
 
@@ -135,6 +138,28 @@ def array_type(self, value):
         else:
             self._array_type = ArrayTypeSchema(value)
 
+    @property
+    def attrs(self) -> AttrsSchema:
+        return self._attrs
+
+    @attrs.setter
+    def attrs(self, value):
+        if value is None or isinstance(value, AttrsSchema):
+            self._attrs = value
+        else:
+            self._attrs = AttrsSchema(value)
+
+    @property
+    def coords(self) -> Mapping:
+        return self._coords
+
+    @coords.setter
+    def coords(self, value):
+        if value is None or isinstance(value, CoordsSchema):
+            self._coords = value
+        else:
+            self._coords = CoordsSchema(value)
+
     @property
     def checks(self) -> List[Callable]:
         return self._checks
@@ -180,14 +205,14 @@ def validate(self, da: xr.DataArray) -> None:
         if self.shape is not None:
             self.shape.validate(da.shape)
 
-        if self.coords is not None:  # pragma: no cover
-            raise NotImplementedError('coords schema not implemented yet')
+        if self.coords is not None:
+            self.coords.validate(da.coords)
 
         if self.chunks is not None:
             self.chunks.validate(da.chunks, da.dims, da.shape)
 
-        if self.attrs:  # pragma: no cover
-            raise NotImplementedError('attrs schema not implemented yet')
+        if self.attrs:
+            self.attrs.validate(da.attrs)
 
         if self.array_type is not None:
             self.array_type.validate(da.data)
@@ -204,3 +229,47 @@ def json(self) -> dict:
             except AttributeError:
                 pass
         return obj
+
+
+class CoordsSchema(BaseSchema):
+
+    _json_schema = {'type': 'string'}
+
+    def __init__(
+        self,
+        coords: Mapping[str, Any],
+        require_all_keys: bool = True,
+        allow_extra_keys: bool = True,
+    ) -> None:
+        self.coords = coords
+        self.require_all_keys = require_all_keys
+        self.allow_extra_keys = allow_extra_keys
+
+    def validate(self, coords: Any) -> None:
+        '''Validate coords
+
+        Parameters
+        ----------
+        coords : dict_like
+            coords of the DataArray. `None` may be used as a wildcard value.
+        '''
+
+        if self.require_all_keys:
+            missing_keys = set(self.coords) - set(coords)
+            if missing_keys:
+                raise SchemaError(f'coords has missing keys: {missing_keys}')
+
+        if not self.allow_extra_keys:
+            extra_keys = set(coords) - set(self.coords)
+            if extra_keys:
+                raise SchemaError(f'coords has extra keys: {extra_keys}')
+
+        for key, da_schema in self.coords.items():
+            if key not in coords:
+                raise SchemaError(f'key {key} not in coords')
+            else:
+                da_schema.validate(coords[key])
+
+    @property
+    def json(self) -> dict:
+        return {k: v.json for k, v in self.attrs.items()}

xarray_schema/dataset.py

@@ -59,16 +59,16 @@ def validate(self, ds: xr.Dataset) -> None:
         if self.coords is not None:  # pragma: no cover
             raise NotImplementedError('coords schema not implemented yet')
 
-        if self.attrs:  # pragma: no cover
-            raise NotImplementedError('attrs schema not implemented yet')
+        if self.attrs:
+            self.attrs.validate(ds.attrs)
 
         if self.checks:
             for check in self.checks:
                 check(ds)
 
     @property
     def json(self):
-        obj = {'data_vars': {}}  # TODO: add when , 'coords': {}, 'attrs': {}}
+        obj = {'data_vars': {}, 'attrs': self.attrs.json if self.attrs is not None else None}
         if self.data_vars:
             for key, var in self.data_vars.items():
                 obj['data_vars'][key] = var.json