Merge pull request #602 from dcs4cop/pont-555-tz-naive-index

Tile server: ensure timezone-naive time selector

Author: pont-us

test/core/test_tile.py

@@ -1,10 +1,12 @@
 import unittest
 
 import numpy as np
+import pandas as pd
 import xarray as xr
 
 from xcube.core.tile import get_var_valid_range
 from xcube.core.tile import parse_non_spatial_labels
+from xcube.core.tile import _ensure_time_compatible
 
 
 class GetVarValidRangeTest(unittest.TestCase):
@@ -102,3 +104,49 @@ def test_invalid(self):
                                      coords=self.coords)
         self.assertEqual("'jetzt' is not a valid value for dimension 'time'",
                          f'{cm.exception}')
+
+    def test_ensure_timezone_naive(self):
+        da_tznaive = xr.DataArray(
+            np.zeros((3,3,3)),
+            coords=self.coords,
+            dims=self.dims)
+        labels = parse_non_spatial_labels(dict(time='2000-01-02T00:00:00Z'),
+                                          dims=da_tznaive.dims,
+                                          coords=da_tznaive.coords,
+                                          var=da_tznaive)
+        self.assertIsNone(pd.Timestamp(labels['time']).tzinfo)
+
+
+class EnsureTimeCompatibleTest(unittest.TestCase):
+    da_tznaive = xr.DataArray(
+        np.arange(1, 4),
+        coords=dict(time=np.arange('2000-01-01', '2000-01-04',
+                                   dtype=np.datetime64)),
+        dims=['time'])
+    da_tzaware = xr.DataArray(
+        np.arange(1, 4),
+        coords=dict(time=pd.date_range("2000-01-01", "2000-01-03", tz='UTC')),
+        dims=['time'])
+    labels_tznaive = dict(time='2000-01-02')
+    labels_tzaware = dict(time='2000-01-02T00:00:00Z')
+
+    def test_both_tznaive(self):
+        self.assertEqual(self.labels_tznaive,
+                         _ensure_time_compatible(self.da_tznaive,
+                                                 self.labels_tznaive))
+
+    def test_both_tzaware(self):
+        self.assertEqual(self.labels_tzaware,
+                         _ensure_time_compatible(self.da_tzaware,
+                                                 self.labels_tzaware))
+
+    def test_tznaive_array_tzaware_indexer(self):
+        self.assertTrue(
+            _are_times_equal(
+                self.labels_tznaive,
+                _ensure_time_compatible(self.da_tznaive,
+                                        self.labels_tzaware)))
+
+
+def _are_times_equal(labels1, labels2):
+    return pd.Timestamp(labels1['time']) == pd.Timestamp(labels2['time'])

xcube/core/tile.py

@@ -250,6 +250,29 @@ def _get_var_2d_array(var: xr.DataArray,
     return array
 
 
+def _ensure_time_compatible(var: xr.DataArray,
+                            labels: Dict[str, Any]) -> Dict[str, Any]:
+    """Ensure that labels['time'] is timezone-naive, if necessary.
+    If var has a 'time' dimension of type datetime64 and labels has a 'time'
+    key with a timezone-aware value, return a modified labels dictionary with
+    a timezone-naive time value. Otherwise return the original labels.
+    """
+    if _has_datetime64_time(var) and \
+       'time' in labels and pd.Timestamp(labels['time']).tzinfo is not None:
+        naive_time = pd.Timestamp(labels['time']).tz_convert(None)
+        return dict(labels, time=naive_time)
+    else:
+        return labels
+
+
+def _has_datetime64_time(var: xr.DataArray) -> bool:
+    """Report whether var has a time dimension with type datetime64"""
+    return 'time' in var.dims and \
+           hasattr(var['time'], 'dtype') and \
+           hasattr(var['time'].dtype, 'type') and \
+           var['time'].dtype.type is np.datetime64
+
+
 def get_var_cmap_params(var: xr.DataArray,
                         cmap_name: Optional[str],
                         cmap_range: Tuple[Optional[float], Optional[float]],
@@ -304,7 +327,8 @@ def parse_non_spatial_labels(
         dims: Sequence[Hashable],
         coords: Mapping[Hashable, xr.DataArray],
         allow_slices: bool = False,
-        exception_type: type = ValueError
+        exception_type: type = ValueError,
+        var: xr.DataArray = None
 ) -> Mapping[str, Any]:
     xy_var_names = get_dataset_xy_var_names(coords, must_exist=False)
     if xy_var_names is None:
@@ -360,4 +384,7 @@ def to_datetime(datetime_str: str, dim_var: xr.DataArray):
             raise exception_type(f'{label_str!r} is not a valid'
                                  f' value for dimension {dim!r}') from e
 
-    return parsed_labels
+    if var is not None:
+        return _ensure_time_compatible(var, parsed_labels)
+    else:
+        return parsed_labels

xcube/webapi/controllers/tiles.py

@@ -79,7 +79,8 @@ def get_dataset_tile(ctx: ServiceContext,
                                       var.dims,
                                       var.coords,
                                       allow_slices=False,
-                                      exception_type=ServiceBadRequestError)
+                                      exception_type=ServiceBadRequestError,
+                                      var=var)
 
     return get_ml_dataset_tile(ml_dataset,
                                var_name,