Force 'dim' to be OrderedDict for python 3.5

Author: pletchm

xarray/core/dataarray.py

@@ -1,4 +1,5 @@
 import functools
+import sys
 import warnings
 from collections import OrderedDict
 
@@ -1153,18 +1154,22 @@ def expand_dims(self, dim=None, axis=None, **dim_kwargs):
             with length 1. If provided as a dict, then the keys are the new
             dimensions and the values are either integers (giving the length of
             the new dimensions) or sequence/ndarray (giving the coordinates of
-            the new dimensions).
+            the new dimensions). **WARNING** for python 3.5, if ``dim`` is
+            dict-like, then it must be an ``OrderedDict`` to ensure that the
+            order in which the dims are given is maintained.
         axis : integer, list (or tuple) of integers, or None
             Axis position(s) where new axis is to be inserted (position(s) on
             the result array). If a list (or tuple) of integers is passed,
             multiple axes are inserted. In this case, dim arguments should be
             same length list. If axis=None is passed, all the axes will be
             inserted to the start of the result array.
-        **dim_kwargs : `int or sequence/ndarray`
+        **dim_kwargs : int or sequence/ndarray
             The keywords are arbitrary dimensions being inserted and the values
             are either the lengths of the new dims (if int is given), or their
             coordinates. Note, this is an alternative to passing a dict to the
-            dim kwarg and will only be used if dim is None.
+            dim kwarg and will only be used if dim is None. **WARNING** in
+            python 3.5, the order in which the dim kwargs are given will not be
+            maintained.
 
         Returns
         -------
@@ -1174,11 +1179,22 @@ def expand_dims(self, dim=None, axis=None, **dim_kwargs):
         if isinstance(dim, int):
             raise TypeError('dim should be str or sequence of strs or dict')
         elif isinstance(dim, str):
-            dim = {dim: 1}
+            dim = OrderedDict(((dim, 1),))
         elif isinstance(dim, (list, tuple)):
             if len(dim) != len(set(dim)):
                 raise ValueError('dims should not contain duplicate values.')
-            dim = {d: 1 for d in dim}
+            dim = OrderedDict(((d, 1) for d in dim))
+
+        # TODO: get rid of the below code block when python 3.5 is no longer
+        #   supported.
+        python36_plus = sys.version_info[0] == 3 and sys.version_info[1] > 5
+        not_ordereddict = dim is not None and not isinstance(dim, OrderedDict)
+        if not python36_plus and not_ordereddict:
+            raise TypeError("dim must be an OrderedDict for python <3.6")
+        # For python 3.5 dim, must be an OrderedDict. This allows kwargs to be
+        # used for python 3.5, but insertion order will not be maintained
+        dim_kwargs = OrderedDict(dim_kwargs)
+
         dim = either_dict_or_kwargs(dim, dim_kwargs, 'expand_dims')
         ds = self._to_temp_dataset().expand_dims(dim, axis)
         return self._from_temp_dataset(ds)

xarray/core/dataset.py

@@ -2343,18 +2343,22 @@ def expand_dims(self, dim=None, axis=None, **dim_kwargs):
             with length 1. If provided as a dict, then the keys are the new
             dimensions and the values are either integers (giving the length of
             the new dimensions) or sequence/ndarray (giving the coordinates of
-            the new dimensions).
+            the new dimensions). **WARNING** for python 3.5, if ``dim`` is
+            dict-like, then it must be an ``OrderedDict`` to ensure that the
+            order in which the dims are given is maintained.
         axis : integer, list (or tuple) of integers, or None
             Axis position(s) where new axis is to be inserted (position(s) on
             the result array). If a list (or tuple) of integers is passed,
             multiple axes are inserted. In this case, dim arguments should be
             same length list. If axis=None is passed, all the axes will be
             inserted to the start of the result array.
-        **dim_kwargs : `int or sequence/ndarray`
+        **dim_kwargs : int or sequence/ndarray
             The keywords are arbitrary dimensions being inserted and the values
             are either the lengths of the new dims (if int is given), or their
             coordinates. Note, this is an alternative to passing a dict to the
-            dim kwarg and will only be used if dim is None.
+            dim kwarg and will only be used if dim is None. **WARNING** in
+            python 3.5, the order in which the dim kwargs are given will not be
+            maintained.
 
         Returns
         -------
@@ -2364,11 +2368,19 @@ def expand_dims(self, dim=None, axis=None, **dim_kwargs):
         if isinstance(dim, int):
             raise TypeError('dim should be str or sequence of strs or dict')
         elif isinstance(dim, str):
-            dim = {dim: 1}
+            dim = OrderedDict(((dim, 1),))
         elif isinstance(dim, (list, tuple)):
             if len(dim) != len(set(dim)):
                 raise ValueError('dims should not contain duplicate values.')
-            dim = {d: 1 for d in dim}
+            dim = OrderedDict(((d, 1) for d in dim))
+
+        # TODO: get rid of the below code block when python 3.5 is no longer
+        #   supported.
+        python36_plus = sys.version_info[0] == 3 and sys.version_info[1] > 5
+        not_ordereddict = dim is not None and not isinstance(dim, OrderedDict)
+        if not python36_plus and not_ordereddict:
+            raise TypeError("dim must be an OrderedDict for python <3.6")
+
         dim = either_dict_or_kwargs(dim, dim_kwargs, 'expand_dims')
 
         if axis is not None and not isinstance(axis, (list, tuple)):

xarray/tests/test_dataarray.py

@@ -1332,11 +1332,11 @@ def test_expand_dims_error(self):
                           coords={'x': np.linspace(0.0, 1.0, 3)},
                           attrs={'key': 'entry'})
         with pytest.raises(TypeError):
-            array.expand_dims({"new_dim": 3.2})
+            array.expand_dims(OrderedDict((("new_dim", 3.2),)))
 
         # Attempt to use both dim and kwargs
         with pytest.raises(ValueError):
-            array.expand_dims({"d": 4}, e=4)
+            array.expand_dims(OrderedDict((("d", 4),)), e=4)
 
     def test_expand_dims(self):
         array = DataArray(np.random.randn(3, 4), dims=['x', 'dim_0'],
@@ -1403,10 +1403,6 @@ def test_expand_dims_with_scalar_coordinate(self):
         assert_identical(array, roundtripped)
 
     def test_expand_dims_with_greater_dim_size(self):
-        """Python 3.6+ dicts keep insertion order, unlike Python 3.5 and
-        earlier. Therefore the following tests have to use ordered dicts to
-        pass for python 3.5 and earlier.
-        """
         array = DataArray(np.random.randn(3, 4), dims=['x', 'dim_0'],
                           coords={'x': np.linspace(0.0, 1.0, 3), 'z': 1.0},
                           attrs={'key': 'entry'})
@@ -1426,19 +1422,15 @@ def test_expand_dims_with_greater_dim_size(self):
         assert_identical(expected, actual)
 
         # Test with kwargs instead of passing dict to dim arg.
-        other_way = array.expand_dims(y=2, z=1, dim_1=['a', 'b', 'c'])
-        # Unfortunately, there is no way to maintain insertion order with
-        # kwargs in python 3.5 and earlier, so for now we have to ensure the
-        # dimensions of the expected result are in the same order as the actual
-        # result to allow the test to pass.
-        other_way_expected_coords = OrderedDict()
-        for dim in other_way.dims:
-            other_way_expected_coords[dim] = expected_coords[dim]
+        other_way = array.expand_dims(dim_1=['a', 'b', 'c'])
+
         other_way_expected = DataArray(
-            array.values * np.ones(list(other_way.shape)),
-            coords=other_way_expected_coords,
-            dims=list(other_way_expected_coords.keys()),
-            attrs={'key': 'entry'}).drop(['y', 'dim_0'])
+            array.values * np.ones([3, 3, 4]),
+            coords={'dim_1': ['a', 'b', 'c'],
+                    'x': np.linspace(0, 1, 3),
+                    'dim_0': range(4), 'z': 1.0},
+            dims=['dim_1', 'x', 'dim_0'],
+            attrs={'key': 'entry'}).drop('dim_0')
         assert_identical(other_way_expected, other_way)
 
     def test_set_index(self):

xarray/tests/test_dataset.py

@@ -2021,10 +2021,10 @@ def test_expand_dims_error(self):
                                    'c': np.linspace(0, 1, 5)},
                            attrs={'key': 'entry'})
         with raises_regex(TypeError, 'value of new dimension'):
-            original.expand_dims({"d": 3.2})
+            original.expand_dims(OrderedDict((("d", 3.2),)))
 
         with raises_regex(ValueError, 'both keyword and positional'):
-            original.expand_dims({"d": 4}, e=4)
+            original.expand_dims(OrderedDict((("d", 4),)), e=4)
 
     def test_expand_dims(self):
         original = Dataset({'x': ('a', np.random.randn(3)),
@@ -2062,7 +2062,8 @@ def test_expand_dims(self):
         # Test expanding one dimension to have size > 1 that doesn't have
         # coordinates, and also expanding another dimension to have size > 1
         # that DOES have coordinates.
-        actual = original.expand_dims({"d": 4, "e": ["l", "m", "n"]})
+        actual = original.expand_dims(
+            OrderedDict((("d", 4), ("e", ["l", "m", "n"]))))
 
         expected = Dataset(
             {'x': xr.DataArray(original['x'].values * np.ones([4, 3, 3]),
@@ -2081,8 +2082,20 @@ def test_expand_dims(self):
         assert_identical(actual, expected)
 
         # Test with kwargs instead of passing dict to dim arg.
-        other_way = original.expand_dims(d=4, e=["l", "m", "n"])
-        assert_identical(expected, other_way)
+        other_way = original.expand_dims(e=["l", "m", "n"])
+        other_way_expected = Dataset(
+            {'x': xr.DataArray(original['x'].values * np.ones([3, 3]),
+                               coords=dict(e=['l', 'm', 'n'],
+                                           a=np.linspace(0, 1, 3)),
+                               dims=['e', 'a']),
+             'y': xr.DataArray(original['y'].values * np.ones([3, 4, 3]),
+                               coords=dict(e=['l', 'm', 'n'],
+                                           b=np.linspace(0, 1, 4),
+                                           a=np.linspace(0, 1, 3)),
+                               dims=['e', 'b', 'a'])},
+            coords={'c': np.linspace(0, 1, 5)},
+            attrs={'key': 'entry'})
+        assert_identical(other_way_expected, other_way)
 
     def test_set_index(self):
         expected = create_test_multiindex()