catch numpy arrays in attrs before converting to dict (#1052)

* catch np.arrays in attrs before converting to dict

* pushed time_check and np_check out to utils.py

* changed utils funtion names

* added numpy scalar test

Author: jsignell

xarray/core/dataarray.py

@@ -22,6 +22,7 @@
                        default_index_coordinate,
                        assert_unique_multiindex_level_names)
 from .formatting import format_item
+from .utils import decode_numpy_dict_values, ensure_us_time_resolution
 
 
 def _infer_coords_and_dims(shape, coords, dims):
@@ -1194,29 +1195,25 @@ def to_dict(self):
         Convert this xarray.DataArray into a dictionary following xarray
         naming conventions.
 
-        Useful for coverting to json.
+        Converts all variables and attributes to native Python objects.
+        Useful for coverting to json. To avoid datetime incompatibility
+        use decode_times=False kwarg in xarrray.open_dataset.
 
         See also
         --------
         xarray.DataArray.from_dict
         """
-        d = {'coords': {}, 'attrs': dict(self.attrs), 'dims': self.dims}
-
-        def time_check(val):
-            # needed because of numpy bug GH#7619
-            if np.issubdtype(val.dtype, np.datetime64):
-                val = val.astype('datetime64[us]')
-            elif np.issubdtype(val.dtype, np.timedelta64):
-                val = val.astype('timedelta64[us]')
-            return val
+        d = {'coords': {}, 'attrs': decode_numpy_dict_values(self.attrs),
+             'dims': self.dims}
 
         for k in self.coords:
-            data = time_check(self[k].values).tolist()
-            d['coords'].update({k: {'data': data,
-                                    'dims': self[k].dims,
-                                    'attrs': dict(self[k].attrs)}})
+            data = ensure_us_time_resolution(self[k].values).tolist()
+            d['coords'].update({
+                k: {'data': data,
+                    'dims': self[k].dims,
+                    'attrs': decode_numpy_dict_values(self[k].attrs)}})
 
-        d.update({'data': time_check(self.values).tolist(),
+        d.update({'data': ensure_us_time_resolution(self.values).tolist(),
                   'name': self.name})
         return d
 

xarray/core/dataset.py

@@ -19,7 +19,8 @@
 from .common import ImplementsDatasetReduce, BaseDataObject
 from .merge import (dataset_update_method, dataset_merge_method,
                     merge_data_and_coords)
-from .utils import Frozen, SortedKeysDict, maybe_wrap_array, hashable
+from .utils import (Frozen, SortedKeysDict, maybe_wrap_array, hashable,
+                    decode_numpy_dict_values, ensure_us_time_resolution)
 from .variable import (Variable, as_variable, IndexVariable, broadcast_variables)
 from .pycompat import (iteritems, basestring, OrderedDict,
                        dask_array_type)
@@ -1917,33 +1918,29 @@ def to_dict(self):
         Convert this dataset to a dictionary following xarray naming
         conventions.
 
-        Useful for coverting to json.
+        Converts all variables and attributes to native Python objects
+        Useful for coverting to json. To avoid datetime incompatibility
+        use decode_times=False kwarg in xarrray.open_dataset.
 
         See also
         --------
         xarray.Dataset.from_dict
         """
-        d = {'coords': {}, 'attrs': dict(self.attrs), 'dims': dict(self.dims),
-             'data_vars': {}}
-
-        def time_check(val):
-            # needed because of numpy bug GH#7619
-            if np.issubdtype(val.dtype, np.datetime64):
-                val = val.astype('datetime64[us]')
-            elif np.issubdtype(val.dtype, np.timedelta64):
-                val = val.astype('timedelta64[us]')
-            return val
+        d = {'coords': {}, 'attrs': decode_numpy_dict_values(self.attrs),
+             'dims': dict(self.dims), 'data_vars': {}}
 
         for k in self.coords:
-            data = time_check(self[k].values).tolist()
-            d['coords'].update({k: {'data': data,
-                                    'dims': self[k].dims,
-                                    'attrs': dict(self[k].attrs)}})
+            data = ensure_us_time_resolution(self[k].values).tolist()
+            d['coords'].update({
+                k: {'data': data,
+                    'dims': self[k].dims,
+                    'attrs': decode_numpy_dict_values(self[k].attrs)}})
         for k in self.data_vars:
-            data = time_check(self[k].values).tolist()
-            d['data_vars'].update({k: {'data': data,
-                                       'dims': self[k].dims,
-                                       'attrs': dict(self[k].attrs)}})
+            data = ensure_us_time_resolution(self[k].values).tolist()
+            d['data_vars'].update({
+                k: {'data': data,
+                    'dims': self[k].dims,
+                    'attrs': decode_numpy_dict_values(self[k].attrs)}})
         return d
 
     @classmethod

xarray/core/utils.py

@@ -460,3 +460,25 @@ def hashable(v):
 
 def not_implemented(*args, **kwargs):
     return NotImplemented
+
+
+def decode_numpy_dict_values(attrs):
+    """Convert attribute values from numpy objects to native Python objects,
+    for use in to_dict"""
+    attrs = dict(attrs)
+    for k, v in attrs.items():
+        if isinstance(v, np.ndarray):
+            attrs[k] = v.tolist()
+        elif isinstance(v, np.generic):
+            attrs[k] = np.asscalar(v)
+    return attrs
+
+
+def ensure_us_time_resolution(val):
+    """Convert val out of numpy time, for use in to_dict.
+    Needed because of numpy bug GH#7619"""
+    if np.issubdtype(val.dtype, np.datetime64):
+        val = val.astype('datetime64[us]')
+    elif np.issubdtype(val.dtype, np.timedelta64):
+        val = val.astype('timedelta64[us]')
+    return val

xarray/test/test_dataarray.py

@@ -1989,6 +1989,24 @@ def test_to_and_from_dict_with_nan_nat(self):
         roundtripped = DataArray.from_dict(da.to_dict())
         self.assertDataArrayIdentical(da, roundtripped)
 
+    def test_to_dict_with_numpy_attrs(self):
+        # this doesn't need to roundtrip
+        x = np.random.randn(10, 3)
+        t = list('abcdefghij')
+        lat = [77.7, 83.2, 76]
+        attrs = {'created': np.float64(1998),
+                 'coords': np.array([37, -110.1, 100]),
+                 'maintainer': 'bar'}
+        da = DataArray(x, {'t': t, 'lat': lat}, dims=['t', 'lat'],
+                       attrs=attrs)
+        expected_attrs = {'created': np.asscalar(attrs['created']),
+                          'coords': attrs['coords'].tolist(),
+                          'maintainer': 'bar'}
+        actual = da.to_dict()
+
+        # check that they are identical
+        self.assertEqual(expected_attrs, actual['attrs'])
+
     def test_to_masked_array(self):
         rs = np.random.RandomState(44)
         x = rs.random_sample(size=(10, 20))

xarray/test/test_dataset.py

@@ -2120,6 +2120,25 @@ def test_to_and_from_dict_with_nan_nat(self):
         roundtripped = Dataset.from_dict(ds.to_dict())
         self.assertDatasetIdentical(ds, roundtripped)
 
+    def test_to_dict_with_numpy_attrs(self):
+        # this doesn't need to roundtrip
+        x = np.random.randn(10)
+        y = np.random.randn(10)
+        t = list('abcdefghij')
+        attrs = {'created': np.float64(1998),
+                 'coords': np.array([37, -110.1, 100]),
+                 'maintainer': 'bar'}
+        ds = Dataset(OrderedDict([('a', ('t', x, attrs)),
+                                  ('b', ('t', y, attrs)),
+                                  ('t', ('t', t))]))
+        expected_attrs = {'created': np.asscalar(attrs['created']),
+                          'coords': attrs['coords'].tolist(),
+                          'maintainer': 'bar'}
+        actual = ds.to_dict()
+
+        # check that they are identical
+        self.assertEqual(expected_attrs, actual['data_vars']['a']['attrs'])
+
     def test_pickle(self):
         data = create_test_data()
         roundtripped = pickle.loads(pickle.dumps(data))