Encode/decode time information to cdms2

Author: shoyer

xray/convert.py

@@ -3,30 +3,26 @@
 import numpy as np
 
 from .core.dataarray import DataArray
-
+from .conventions import (
+    maybe_encode_timedelta, maybe_encode_datetime, decode_cf)
 
 ignored_attrs = set(['name', 'tileIndex'])
 
 
-def _get_cdms2_attrs(var):
-    return dict((k, v) for k, v in var.attributes.items()
-                if k not in ignored_attrs)
-
-
 def from_cdms2(variable):
     """Convert a cdms2 variable into an DataArray
     """
+    def get_cdms2_attrs(var):
+        return dict((k, v) for k, v in var.attributes.items()
+                    if k not in ignored_attrs)
+
     values = np.asarray(variable)
     name = variable.id
-    coords = [(v.id, np.asarray(v), _get_cdms2_attrs(v))
+    coords = [(v.id, np.asarray(v), get_cdms2_attrs(v))
               for v in variable.getAxisList()]
-    attrs = _get_cdms2_attrs(variable)
-    return DataArray(values, coords=coords, name=name, attrs=attrs)
-
-
-def _set_cdms2_attrs(var, attrs):
-    for k, v in attrs.items():
-        setattr(var, k, v)
+    attrs = get_cdms2_attrs(variable)
+    dataarray = DataArray(values, coords=coords, name=name, attrs=attrs)
+    return decode_cf(dataarray.to_dataset())[dataarray.name]
 
 
 def to_cdms2(dataarray):
@@ -35,13 +31,21 @@ def to_cdms2(dataarray):
     # we don't want cdms2 to be a hard dependency
     import cdms2
 
+    def encode(var):
+        return maybe_encode_timedelta(maybe_encode_datetime(var))
+
+    def set_cdms2_attrs(var, attrs):
+        for k, v in attrs.items():
+            setattr(var, k, v)
+
     axes = []
     for dim in dataarray.dims:
-        coord = dataarray.coords[dim]
+        coord = encode(dataarray.coords[dim])
         axis = cdms2.createAxis(coord.values, id=dim)
-        _set_cdms2_attrs(axis, coord.attrs)
+        set_cdms2_attrs(axis, coord.attrs)
         axes.append(axis)
 
-    var = cdms2.createVariable(dataarray.values, axes=axes, id=dataarray.name)
-    _set_cdms2_attrs(var, dataarray.attrs)
-    return var
+    var = encode(dataarray)
+    cdms2_var = cdms2.createVariable(var.values, axes=axes, id=dataarray.name)
+    set_cdms2_attrs(cdms2_var, var.attrs)
+    return cdms2_var

xray/test/test_dataarray.py

@@ -1002,18 +1002,24 @@ def test_to_and_from_cdms2(self):
             raise unittest.SkipTest('cdms2 not installed')
 
         original = DataArray(np.arange(6).reshape(2, 3),
-                             [('xxx', [-2, 2], {'units': 'meters'}),
-                              ('yyy', [3, 4, 5])],
+                             [('distance', [-2, 2], {'units': 'meters'}),
+                              ('time', pd.date_range('2000-01-01', periods=3))],
                              name='foo', attrs={'baz': 123})
+        expected_coords = [Coordinate('distance', [-2, 2]),
+                           Coordinate('time', [0, 1, 2])]
         actual = original.to_cdms2()
         self.assertArrayEqual(actual, original)
         self.assertEqual(actual.id, original.name)
         self.assertItemsEqual(actual.getAxisIds(), original.dims)
-        for axis, coord in zip(actual.getAxisList(), original.coords.values()):
+        for axis, coord in zip(actual.getAxisList(), expected_coords):
             self.assertEqual(axis.id, coord.name)
-            self.assertArrayEqual(axis, coord)
+            self.assertArrayEqual(axis, coord.values)
         self.assertEqual(actual.baz, original.attrs['baz'])
 
+        component_times = actual.getAxis(1).asComponentTime()
+        self.assertEqual(len(component_times), 3)
+        self.assertEqual(str(component_times[0]), '2000-1-1 0:0:0.0')
+
         roundtripped = DataArray.from_cdms2(actual)
         self.assertDataArrayIdentical(original, roundtripped)