check chunks

Author: Joseph Hamman

tests/test_core.py

@@ -81,6 +81,28 @@ def test_dataarray_validate_array_type():
         schema.validate(da)
 
 
+def test_dataarray_validate_chunks():
+    pytest.importorskip('dask')
+
+    da = xr.DataArray(np.ones(4), dims=['x']).chunk({'x': 2})
+    schema = DataArraySchema(chunks={'x': 2})
+    schema.validate(da)
+
+    schema = DataArraySchema(chunks={'x': (2, 2)})
+    schema.validate(da)
+
+    schema = DataArraySchema(chunks={'x': [2, 2]})
+    schema.validate(da)
+
+    schema = DataArraySchema(chunks={'x': 3})
+    with pytest.raises(SchemaError, match=r'.*(3).*'):
+        schema.validate(da)
+
+    schema = DataArraySchema(chunks={'x': (2, 1)})
+    with pytest.raises(SchemaError, match=r'.*(2, 1).*'):
+        schema.validate(da)
+
+
 def test_dataset_empty_constructor():
     ds_schema = DatasetSchema()
     assert hasattr(ds_schema, 'validate')

xarray_schema/core.py

@@ -103,7 +103,17 @@ def validate(self, da: xr.DataArray) -> xr.DataArray:
             raise NotImplementedError('coords schema not implemented yet')
 
         if self.chunks:
-            raise NotImplementedError('chunk schema not implemented yet')
+            dim_chunks = dict(zip(da.dims, da.chunks))
+            for key, ec in self.chunks.items():
+                if isinstance(ec, int):
+                    for ac in dim_chunks[key][:-1]:
+                        if ac != ec:
+                            raise SchemaError(f'{key} chunks did not match: {ac} != {ec}')
+
+                else:  # assumes ec is an iterable
+                    ac = dim_chunks[key]
+                    if tuple(ac) != tuple(ec):
+                        raise SchemaError(f'{key} chunks did not match: {ac} != {ec}')
 
         if self.attrs:
             raise NotImplementedError('attrs schema not implemented yet')