Cleanup xarray provider (#2224)

* Cleanup xarray provider

* Respond to PR feedback

* Fix xarray tests

Author: webb-ben

pygeoapi/provider/xarray_.py

@@ -126,6 +126,11 @@ def get_fields(self):
                     elif dtype.name.startswith('str'):
                         dtype = 'string'
 
+                    if value.attrs.get('units') is None:
+                        msg = f'Field {key} missing units, will be skipped'
+                        LOGGER.warning(msg)
+                        continue
+
                     self._fields[key] = {
                         'type': dtype,
                         'title': value.attrs.get('long_name'),
@@ -249,19 +254,21 @@ def query(self, properties=[], subsets={}, bbox=[], bbox_crs=4326,
                 data.coords[self.x_field].values[-1],
                 data.coords[self.y_field].values[-1]
             ],
-            "driver": "xarray",
-            "height": data.sizes[self.y_field],
-            "width": data.sizes[self.x_field],
-            "variables": {var_name: var.attrs
-                          for var_name, var in data.variables.items()}
+            'driver': 'xarray',
+            'height': data.sizes[self.y_field],
+            'width': data.sizes[self.x_field],
+            'variables': {
+                var_name: var.attrs
+                for var_name, var in data.variables.items()
+            }
         }
 
         if self.time_field is not None:
             out_meta['time'] = [
                 _to_datetime_string(data.coords[self.time_field].values[0]),
-                _to_datetime_string(data.coords[self.time_field].values[-1]),
+                _to_datetime_string(data.coords[self.time_field].values[-1])
             ]
-            out_meta["time_steps"] = data.sizes[self.time_field]
+            out_meta['time_steps'] = data.sizes[self.time_field]
 
         LOGGER.debug('Serializing data in memory')
         if format_ == 'json':
@@ -395,25 +402,30 @@ def gen_covjson(self, metadata, data, fields):
         try:
             for key, value in selected_fields.items():
                 LOGGER.debug(f'Adding range {key}')
-                cj['ranges'][key] = {
+                range = {
                     'type': 'NdArray',
                     'dataType': value['type'],
                     'axisNames': [
                         'y', 'x'
                     ],
-                    'shape': [metadata['height'],
-                              metadata['width']]
+                    'shape': [
+                        metadata['height'], metadata['width']
+                    ],
+                    'values': [
+                        None if np.isnan(v) else v
+                        for v in data[key].values.flatten()
+                    ]
                 }
-                cj['ranges'][key]['values'] = [
-                    None if np.isnan(v) else v
-                    for v in data[key].values.flatten()
-                ]
 
                 if self.time_field is not None:
-                    cj['ranges'][key]['axisNames'].append('t')
-                    cj['ranges'][key]['shape'].append(metadata['time_steps'])
+                    LOGGER.debug(f'Adding time axis to range {key}')
+                    range['axisNames'].insert(0, 't')
+                    range['shape'].insert(0, metadata['time_steps'])
+
+                cj['ranges'][key] = range
+
         except IndexError as err:
-            LOGGER.warning(err)
+            LOGGER.error(err)
             raise ProviderQueryError('Invalid query parameter')
 
         LOGGER.debug('Returning data')
@@ -684,11 +696,11 @@ def _get_zarr_data(data):
 
 def _convert_float32_to_float64(data):
     """
-        Converts DataArray values of float32 to float64
-        :param data: Xarray dataset of coverage data
+    Converts DataArray values of float32 to float64
+    :param data: Xarray dataset of coverage data
 
-        :returns: Xarray dataset of coverage data
-        """
+    :returns: Xarray dataset of coverage data
+    """
 
     for var_name in data.variables:
         if data[var_name].dtype == 'float32':

pygeoapi/provider/xarray_edr.py

@@ -36,7 +36,7 @@
 from pygeoapi.provider.xarray_ import (
     _to_datetime_string,
     _convert_float32_to_float64,
-    XarrayProvider,
+    XarrayProvider
 )
 
 LOGGER = logging.getLogger(__name__)
@@ -73,10 +73,9 @@ def position(self, **kwargs):
 
         query_params = {}
 
+        LOGGER.debug('Query type: position')
         LOGGER.debug(f'Query parameters: {kwargs}')
 
-        LOGGER.debug(f"Query type: {kwargs.get('query_type')}")
-
         wkt = kwargs.get('wkt')
         if wkt is not None:
             LOGGER.debug('Processing WKT')
@@ -115,7 +114,10 @@ def position(self, **kwargs):
 
         try:
             if select_properties:
-                self._fields = {k: v for k, v in self._fields.items() if k in select_properties}  # noqa
+                self._fields = {
+                    k: v for k, v in self._fields.items()
+                    if k in select_properties
+                }
                 data = self._data[[*select_properties]]
             else:
                 data = self._data
@@ -156,12 +158,12 @@ def position(self, **kwargs):
         bbox = wkt.bounds
         out_meta = {
             'bbox': [bbox[0], bbox[1], bbox[2], bbox[3]],
-            "time": time,
-            "driver": "xarray",
-            "height": height,
-            "width": width,
-            "time_steps": time_steps,
-            "variables": {var_name: var.attrs
+            'time': time,
+            'driver': 'xarray',
+            'height': height,
+            'width': width,
+            'time_steps': time_steps,
+            'variables': {var_name: var.attrs
                           for var_name, var in data.variables.items()}
         }
 
@@ -183,12 +185,11 @@ def cube(self, **kwargs):
 
         query_params = {}
 
+        LOGGER.debug('Query type: cube')
         LOGGER.debug(f'Query parameters: {kwargs}')
 
-        LOGGER.debug(f"Query type: {kwargs.get('query_type')}")
-
         bbox = kwargs.get('bbox')
-        xmin, ymin, xmax, ymax = self._configure_bbox(bbox)
+        xmin, ymin, xmax, ymax = self._configure_bbox()
 
         if len(bbox) == 4:
             query_params[self.x_field] = slice(bbox[xmin], bbox[xmax])
@@ -208,15 +209,17 @@ def cube(self, **kwargs):
         if datetime_ is not None:
             query_params[self.time_field] = self._make_datetime(datetime_)
 
+        fields = {
+            field: self.fields[field]
+            for field in select_properties
+            if field in self.fields
+        } if select_properties else self.fields
+
         LOGGER.debug(f'query parameters: {query_params}')
         try:
-            if select_properties:
-                self._fields = {k: v for k, v in self._fields.items() if k in select_properties}  # noqa
-                data = self._data[[*select_properties]]
-            else:
-                data = self._data
-            data = data.sel(query_params)
-            data = _convert_float32_to_float64(data)
+            data = _convert_float32_to_float64(
+                self._data[[*fields]].sel(query_params)
+            )
         except KeyError:
             raise ProviderNoDataError()
 
@@ -231,16 +234,18 @@ def cube(self, **kwargs):
                 data.coords[self.x_field].values[-1],
                 data.coords[self.y_field].values[-1]
             ],
-            "time": time,
-            "driver": "xarray",
-            "height": height,
-            "width": width,
-            "time_steps": time_steps,
-            "variables": {var_name: var.attrs
-                          for var_name, var in data.variables.items()}
+            'time': time,
+            'driver': 'xarray',
+            'height': height,
+            'width': width,
+            'time_steps': time_steps,
+            'variables': {
+                var_name: var.attrs
+                for var_name, var in data.variables.items()
+            }
         }
 
-        return self.gen_covjson(out_meta, data, self.fields)
+        return self.gen_covjson(out_meta, data, fields)
 
     def _make_datetime(self, datetime_):
         """
@@ -300,7 +305,7 @@ def _parse_time_metadata(self, data, kwargs):
             time_steps = kwargs.get('limit')
         return time, time_steps
 
-    def _configure_bbox(self, bbox):
+    def _configure_bbox(self):
         xmin, ymin, xmax, ymax = 0, 1, 2, 3
         if self._data[self.x_field][0] > self._data[self.x_field][-1]:
             xmin, xmax = xmax, xmin

tests/provider/test_xarray_zarr_provider.py

@@ -73,7 +73,7 @@ def config_no_time(tmp_path):
 def test_provider(config):
     p = XarrayProvider(config)
 
-    assert len(p.fields) == 4
+    assert len(p.fields) == 3
     assert len(p.axes) == 3
     assert p.axes == ['lon', 'lat', 'time']
 
@@ -82,7 +82,7 @@ def test_schema(config):
     p = XarrayProvider(config)
 
     assert isinstance(p.fields, dict)
-    assert len(p.fields) == 4
+    assert len(p.fields) == 3
     assert p.fields['analysed_sst']['title'] == 'analysed sea surface temperature'  # noqa
 
 
@@ -107,7 +107,7 @@ def test_numpy_json_serial():
 def test_no_time(config_no_time):
     p = XarrayProvider(config_no_time)
 
-    assert len(p.fields) == 4
+    assert len(p.fields) == 3
     assert p.axes == ['lon', 'lat']
 
     coverage = p.query(format='json')