Checking args, added tests

Author: forman

test/core/gridmapping/test_cfconv.py

@@ -204,24 +204,35 @@ def _gen_2d(self):
 class AddSpatialRefTest(S3Test):
 
     def test_add_spatial_ref(self):
-        original_dataset = new_cube(x_name='x',
-                                    y_name='y',
-                                    x_start=0,
-                                    y_start=0,
-                                    x_res=10,
-                                    y_res=10,
-                                    x_units='metres',
-                                    y_units='metres',
-                                    drop_bounds=True,
-                                    width=100,
-                                    height=100,
-                                    variables=dict(A=1.3, B=8.3))
+        self.assert_add_spatial_ref_ok(None, None)
+        self.assert_add_spatial_ref_ok(None, ('cx', 'cy'))
+        self.assert_add_spatial_ref_ok('crs', None)
+        self.assert_add_spatial_ref_ok('crs', ('cx', 'cy'))
+
+    def assert_add_spatial_ref_ok(self, crs_var_name, xy_dim_names):
 
         root = 'eurodatacube-test/xcube-eea'
         data_id = 'test.zarr'
-        crs_var_name = 'spatial_ref'
         crs = pyproj.CRS.from_string("EPSG:3035")
 
+        if xy_dim_names:
+            x_name, y_name = xy_dim_names
+        else:
+            x_name, y_name = 'x', 'y'
+
+        cube = new_cube(x_name=x_name,
+                        y_name=y_name,
+                        x_start=0,
+                        y_start=0,
+                        x_res=10,
+                        y_res=10,
+                        x_units='metres',
+                        y_units='metres',
+                        drop_bounds=True,
+                        width=100,
+                        height=100,
+                        variables=dict(A=1.3, B=8.3))
+
         storage_options = dict(
             anon=False,
             client_kwargs=dict(
@@ -239,13 +250,13 @@ def test_add_spatial_ref(self):
                                        root=root,
                                        storage_options=storage_options)
 
-        data_store.write_data(original_dataset, data_id=data_id)
-        opened_dataset = data_store.open_data(data_id)
-        self.assertEqual({'A', 'B', 'x', 'y', 'time'},
-                         set(opened_dataset.variables))
+        data_store.write_data(cube, data_id=data_id)
+        cube = data_store.open_data(data_id)
+        self.assertEqual({'A', 'B', 'time', x_name, y_name},
+                         set(cube.variables))
 
         with self.assertRaises(ValueError) as cm:
-            GridMapping.from_dataset(opened_dataset)
+            GridMapping.from_dataset(cube)
         self.assertEqual(
             ('cannot find any grid mapping in dataset',),
             cm.exception.args
@@ -254,21 +265,32 @@ def test_add_spatial_ref(self):
         path = f"{root}/{data_id}"
         group_store = fs.get_mapper(path, create=True)
 
-        add_spatial_ref(group_store,
-                        crs,
-                        crs_var_name=crs_var_name)
-
-        self.assertTrue(fs.exists(f"{path}/{crs_var_name}"))
-        self.assertTrue(fs.exists(f"{path}/{crs_var_name}/.zarray"))
-        self.assertTrue(fs.exists(f"{path}/{crs_var_name}/.zattrs"))
-
-        opened_dataset = data_store.open_data(data_id)
-        self.assertEqual({'A', 'B', 'x', 'y', 'time', 'spatial_ref'},
-                         set(opened_dataset.variables))
-        self.assertEqual(crs_var_name,
-                         opened_dataset.A.attrs.get('grid_mapping'))
-        self.assertEqual(crs_var_name,
-                         opened_dataset.B.attrs.get('grid_mapping'))
-
-        gm = GridMapping.from_dataset(opened_dataset)
+        expected_crs_var_name = crs_var_name or 'spatial_ref'
+
+        self.assertTrue(fs.exists(path))
+        self.assertFalse(fs.exists(f"{path}/{expected_crs_var_name}"))
+        self.assertFalse(fs.exists(f"{path}/{expected_crs_var_name}/.zarray"))
+        self.assertFalse(fs.exists(f"{path}/{expected_crs_var_name}/.zattrs"))
+
+        kwargs = {}
+        if crs_var_name is not None:
+            kwargs.update(crs_var_name=crs_var_name)
+        if xy_dim_names is not None:
+            kwargs.update(xy_dim_names=xy_dim_names)
+        add_spatial_ref(group_store, crs, **kwargs)
+
+        self.assertTrue(fs.exists(f"{path}/{expected_crs_var_name}"))
+        self.assertTrue(fs.exists(f"{path}/{expected_crs_var_name}/.zarray"))
+        self.assertTrue(fs.exists(f"{path}/{expected_crs_var_name}/.zattrs"))
+
+        cube = data_store.open_data(data_id)
+        self.assertEqual({'A', 'B', 'time',
+                          x_name, y_name, expected_crs_var_name},
+                         set(cube.variables))
+        self.assertEqual(expected_crs_var_name,
+                         cube.A.attrs.get('grid_mapping'))
+        self.assertEqual(expected_crs_var_name,
+                         cube.B.attrs.get('grid_mapping'))
+
+        gm = GridMapping.from_dataset(cube)
         self.assertIn("LAEA Europe", gm.crs.srs)

xcube/core/gridmapping/cfconv.py

@@ -20,6 +20,7 @@
 # SOFTWARE.
 
 import warnings
+from collections.abc import MutableMapping
 from typing import Optional, Dict, Any, Hashable, Union, Set, List, Tuple
 
 import numpy as np
@@ -29,6 +30,7 @@
 import zarr.convenience
 
 from xcube.core.schema import get_dataset_chunks
+from xcube.util.assertions import assert_instance
 
 
 class GridCoords:
@@ -290,36 +292,41 @@ def _find_dataset_tile_size(dataset: xr.Dataset,
     return None
 
 
-def add_spatial_ref(group_store: zarr.convenience.StoreLike,
+def add_spatial_ref(dataset_store: zarr.convenience.StoreLike,
                     crs: pyproj.CRS,
                     crs_var_name: Optional[str] = 'spatial_ref',
                     xy_dim_names: Optional[Tuple[str, str]] = None):
     """
     Helper function that allows adding a spatial reference to an
     existing Zarr dataset.
 
-    :param group_store: The dataset's existing Zarr store or path.
-    :param crs: The coordinate reference system.
+    :param dataset_store: The dataset's existing Zarr store or path.
+    :param crs: The spatial coordinate reference system.
     :param crs_var_name: The name of the variable that will hold the
         spatial reference. Defaults to "spatial_ref".
     :param xy_dim_names: The names of the x and y dimensions.
         Defaults to ("x", "y").
     """
+    assert_instance(dataset_store, (MutableMapping, str), name='group_store')
+    assert_instance(crs_var_name, str, name='crs_var_name')
+    x_dim_name, y_dim_name = xy_dim_names or ('x', 'y')
+
     spatial_attrs = crs.to_cf()
     spatial_attrs['_ARRAY_DIMENSIONS'] = []  # Required by xarray
-    group = zarr.open(group_store, mode='r+')
+    group = zarr.open(dataset_store, mode='r+')
     spatial_ref = group.array(crs_var_name,
                               0,
                               shape=(),
                               dtype=np.uint8,
                               fill_value=0)
     spatial_ref.attrs.update(**spatial_attrs)
 
-    yx_dim_names = list(reversed(xy_dim_names or ('x', 'y')))
     for item_name, item in group.items():
         if item_name != crs_var_name:
             dims = item.attrs.get('_ARRAY_DIMENSIONS')
-            if dims and len(dims) >= 2 and dims[-2:] == yx_dim_names:
+            if dims and len(dims) >= 2 \
+                    and dims[-2] == y_dim_name \
+                    and dims[-1] == x_dim_name:
                 item.attrs['grid_mapping'] = crs_var_name
 
-    zarr.convenience.consolidate_metadata(group_store)
+    zarr.convenience.consolidate_metadata(dataset_store)