Adding tests for chunks validation

Author: orianac

tests/test_core.py

@@ -101,6 +101,26 @@ def test_dataarray_validate_chunks():
     schema = DataArraySchema(chunks={'x': (2, 1)})
     with pytest.raises(SchemaError, match=r'.*(2, 1).*'):
         schema.validate(da)
+    
+    # check that when expected chunk == -1 it fails 
+    schema = DataArraySchema(chunks={'x': -1})
+    with pytest.raises(SchemaError, match=r'.*(4).*'):
+        schema.validate(da)
+
+    # check that when chunking schema is -1 it also works
+    # both when chunking is specified as -1 and as 4
+    schema = DataArraySchema(chunks={'x': 4})
+    da = xr.DataArray(np.ones(4), dims=['x']).chunk({'x': -1})
+    schema.validate(da)
+
+    schema = DataArraySchema(chunks={'x': -1})
+    da = xr.DataArray(np.ones(4), dims=['x']).chunk({'x': 4})
+    schema.validate(da)
+
+    schema = DataArraySchema(chunks={'x': -1})
+    da = xr.DataArray(np.ones(4), dims=['x']).chunk({'x': -1})
+    schema.validate(da)
+
 
 
 def test_dataset_empty_constructor():
@@ -125,3 +145,9 @@ def test_dataset_example():
         }
     )
     ds_schema.validate(ds)
+
+def test_validate():
+    schema = DataArraySchema()
+    da = xr.DataArray(np.ones(4), dims=['x']).chunk({'x': (1,2,1)})
+    with pytest.raises(AssertionError, match=r'.*(gracious).*'):
+        schema.validate(da)

xarray_schema/core.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import xarray as xr
-
+from dask.array.core import _check_regular_chunks
 # TODOs:
 # - api grouping, should the constructors look similar to the DataArray/Dataset constructors
 
@@ -72,6 +72,8 @@ def validate(self, da: xr.DataArray) -> xr.DataArray:
         ------
         SchemaError
         '''
+        if da.chunks:
+            assert _check_regular_chunks(da.chunks), 'Good gracious no! Chunks are not regular!'
         if not isinstance(da, xr.DataArray):
             raise ValueError('Input must be a xarray.DataArray')
 
@@ -106,11 +108,12 @@ def validate(self, da: xr.DataArray) -> xr.DataArray:
 
         if self.chunks:
             dim_chunks = dict(zip(da.dims, da.chunks))
+            dim_sizes = dict(zip(da.dims, da.shape))
             for key, ec in self.chunks.items():
                 if isinstance(ec, int):
                     # handles case of expected chunksize is shorthand of -1 which translates to the full length of dimension
-                    if ec==-1:
-                        ec = len(da[key])
+                    if ec == -1:
+                        ec = dim_sizes[key]
                         # grab the first entry in da's tuple of chunks to be representative (as it should be assuming they're regular)
                     ac = dim_chunks[key][0]
                     if ac != ec: