Add rename_vars and rename_dims

xarray.code-workspace

@@ -0,0 +1,11 @@
+{
+	"folders": [
+		{
+			"path": "/Users/jkent/Documents/esmlab-regrid/xarray/xarray"
+		},
+		{
+			"path": "."
+		}
+	],
+	"settings": {}
+}

xarray/core/dataset.py

@@ -149,8 +149,8 @@ def merge_indexes(
 
         for n in var_names:
             var = variables[n]
-            if (current_index_variable is not None and
-                    var.dims != current_index_variable.dims):
+            if (current_index_variable is not None
+                    and var.dims != current_index_variable.dims):
                 raise ValueError(
                     "dimension mismatch between %r %s and %r %s"
                     % (dim, current_index_variable.dims, n, var.dims))
@@ -209,8 +209,8 @@ def split_indexes(
     Not public API. Used in Dataset and DataArray reset_index
     methods.
     """
-    if (isinstance(dims_or_levels, str)
-            or not isinstance(dims_or_levels, Sequence)):
+    if (isinstance(dims_or_levels, str) or
+            not isinstance(dims_or_levels, Sequence)):
         dims_or_levels = [dims_or_levels]
 
     dim_levels \
@@ -287,8 +287,8 @@ def __len__(self) -> int:
         return len(self._dataset._variables) - len(self._dataset._coord_names)
 
     def __contains__(self, key) -> bool:
-        return (key in self._dataset._variables and
-                key not in self._dataset._coord_names)
+        return (key in self._dataset._variables
+                and key not in self._dataset._coord_names)
 
     def __getitem__(self, key) -> 'DataArray':
         if key not in self._dataset._coord_names:
@@ -1188,8 +1188,8 @@ def identical(self, other):
         Dataset.equals
         """
         try:
-            return (utils.dict_equiv(self.attrs, other.attrs) and
-                    self._all_compat(other, 'identical'))
+            return (utils.dict_equiv(self.attrs, other.attrs)
+                    and self._all_compat(other, 'identical'))
         except (TypeError, AttributeError):
             return False
 
@@ -2147,8 +2147,8 @@ def _validate_interp_indexer(x, new_x):
             # In the case of datetimes, the restrictions placed on indexers
             # used with interp are stronger than those which are placed on
             # isel, so we need an additional check after _validate_indexers.
-            if (_contains_datetime_like_objects(x) and
-                    not _contains_datetime_like_objects(new_x)):
+            if (_contains_datetime_like_objects(x)
+                    and not _contains_datetime_like_objects(new_x)):
                 raise TypeError('When interpolating over a datetime-like '
                                 'coordinate, the coordinates to '
                                 'interpolate to must be either datetime '
@@ -2256,23 +2256,45 @@ def interp_like(self, other, method='linear', assume_sorted=False,
         return ds.interp(numeric_coords, method, assume_sorted, kwargs)
 
     # Helper methods for rename()
-    def _rename_vars(self, name_dict, dims_dict):
+    def _rename_var_dims_helper(name_dict, v):
+        dims = tuple(name_dict.get(dim, dim) for dim in v.dims)
+        var = v.copy(deep=False)
+        var.dims = dims
+        return var
+
+    def _rename_var_dims_only(self, name_dict):
+        variables = OrderedDict()
+        for k, v in self.variables.items():
+            variables[k] = _rename_var_dims_helper(name_dict, v)
+        return variables
+
+    def _rename_vars_helper(name_dict, var):
+        name = name_dict.get(k, k)
+        if name in variables:
+            raise ValueError('the new name %r conflicts' % (name,))
+        variables[name] = var
+        if k in self._coord_names:
+            coord_names.add(name)
+        return variables, coord_names
+
+    def _rename_vars_only(self, name_dict):
         variables = OrderedDict()
         coord_names = set()
         for k, v in self.variables.items():
-            name = name_dict.get(k, k)
-            dims = tuple(dims_dict.get(dim, dim) for dim in v.dims)
             var = v.copy(deep=False)
-            var.dims = dims
-            if name in variables:
-                raise ValueError('the new name %r conflicts' % (name,))
-            variables[name] = var
-            if k in self._coord_names:
-                coord_names.add(name)
+            variables, coord_names = _rename_vars_helper(name_dict, var)
+        return variables, coord_names
+
+    def _rename_vars(self, name_dict):
+        variables = OrderedDict()
+        coord_names = set()
+        for k, v in self.variables.items():
+            var = _rename_var_dims_helper(name_dict, v)
+            variables, coord_names = _rename_vars_helper(name_dict, var)
         return variables, coord_names
 
-    def _rename_dims(self, dims_dict):
-        return {dims_dict.get(k, k): v for k, v in self.dims.items()}
+    def _rename_dims(self, name_dict):
+        return {name_dict.get(k, k): v for k, v in self.dims.items()}
 
     def _rename_indexes(self, name_dict):
         if self._indexes is None:
@@ -2289,23 +2311,20 @@ def _rename_indexes(self, name_dict):
             indexes[new_name] = index
         return indexes
 
-    def _rename_all(self, name_dict, dim_dict):
-        variables, coord_names = self._rename_vars(name_dict, dim_dict)
-        dims = self._rename_dims(dim_dict)
+    def _rename_all(self, name_dict):
+        variables, coord_names = self._rename_vars(name_dict)
+        dims = self._rename_dims(name_dict)
         indexes = self._rename_indexes(name_dict)
         return variables, coord_names, dims, indexes
 
-    def rename(self, name_dict=None, inplace=None, **names):
+    def rename(self, name_dict=None, **names):
         """Returns a new object with renamed variables and dimensions.
 
         Parameters
         ----------
         name_dict : dict-like, optional
             Dictionary whose keys are current variable or dimension names and
             whose values are the desired names.
-        inplace : bool, optional
-            If True, rename variables and dimensions in-place. Otherwise,
-            return a new dataset object.
         **names, optional
             Keyword form of ``name_dict``.
             One of name_dict or names must be provided.
@@ -2318,20 +2337,87 @@ def rename(self, name_dict=None, inplace=None, **names):
         See Also
         --------
         Dataset.swap_dims
+        Dataset.rename_vars
+        Dataset.rename_dims
         DataArray.rename
         """
-        # TODO: add separate rename_vars and rename_dims methods.
-        inplace = _check_inplace(inplace)
         name_dict = either_dict_or_kwargs(name_dict, names, 'rename')
-        for k, v in name_dict.items():
+        for  in name_dict.items():
             if k not in self and k not in self.dims:
                 raise ValueError("cannot rename %r because it is not a "
                                  "variable or dimension in this dataset" % k)
 
         variables, coord_names, dims, indexes = self._rename_all(
-            name_dict=name_dict, dim_dict=name_dict)
+            name_dict=name_dict)
         return self._replace(variables, coord_names, dims=dims,
-                             indexes=indexes, inplace=inplace)
+                             indexes=indexes)
+
+    def rename_dims(self, name_dict=None, **names):
+        """Returns a new object with renamed dimensions only.
+
+        Parameters
+        ----------
+        name_dict : dict-like, optional
+            Dictionary whose keys are current dimension names and
+            whose values are the desired names.
+        **names, optional
+            Keyword form of ``name_dict``.
+            One of name_dict or names must be provided.
+
+        Returns
+        -------
+        renamed : Dataset
+            Dataset with renamed dimensions.
+
+        See Also
+        --------
+        Dataset.swap_dims
+        Dataset.rename
+        Dataset.rename_vars
+        DataArray.rename
+        """
+        name_dict = either_dict_or_kwargs(name_dict, names, 'rename')
+        for k in name_dict.items():
+            if k not in self.dims:
+                raise ValueError("cannot rename %r because it is not a "
+                                 "dimension in this dataset" % k)
+
+        dims = self._rename_dims(name_dict)
+        variables = self._rename_var_dims_only(name_dict)
+        indexes = self._rename_indexes(name_dict)
+        return self._replace(variables, dims=dims, indexes=indexes)
+
+    def rename_vars(self, name_dict=None, **names):
+        """Returns a new object with renamed variables including coordinates
+
+        Parameters
+        ----------
+        name_dict : dict-like, optional
+            Dictionary whose keys are current variable or coordinate names and
+            whose values are the desired names.
+        **names, optional
+            Keyword form of ``name_dict``.
+            One of name_dict or names must be provided.
+
+        Returns
+        -------
+        renamed : Dataset
+            Dataset with renamed variables including coordinates
+
+        See Also
+        --------
+        Dataset.swap_dims
+        Dataset.rename
+        Dataset.rename_dims
+        DataArray.rename
+        """
+        name_dict = either_dict_or_kwargs(name_dict, names, 'rename')
+        for k in name_dict.items():
+            if k not in self:
+                raise ValueError("cannot rename %r because it is not a "
+                                 "variable or coordinate in this dataset" % k)
+        variables, coord_names = self._rename_vars_only(name_dict)
+        return self._replace(variables, coord_names)
 
     def swap_dims(self, dims_dict, inplace=None):
         """Returns a new object with swapped dimensions.
@@ -2460,8 +2546,8 @@ def expand_dims(self, dim=None, axis=None, **dim_kwargs):
             if d in self.dims:
                 raise ValueError(
                     'Dimension {dim} already exists.'.format(dim=d))
-            if (d in self._variables and
-                    not utils.is_scalar(self._variables[d])):
+            if (d in self._variables
+                    and not utils.is_scalar(self._variables[d])):
                 raise ValueError(
                     '{dim} already exists as coordinate or'
                     ' variable name.'.format(dim=d))
@@ -3252,9 +3338,9 @@ def reduce(self, func, dim=None, keep_attrs=None, keepdims=False,
                 if not reduce_dims:
                     variables[name] = var
             else:
-                if (not numeric_only or
-                        np.issubdtype(var.dtype, np.number) or
-                        (var.dtype == np.bool_)):
+                if (not numeric_only
+                        or np.issubdtype(var.dtype, np.number)
+                        or (var.dtype == np.bool_)):
                     if len(reduce_dims) == 1:
                         # unpack dimensions for the benefit of functions
                         # like np.argmin which can't handle tuple arguments
@@ -3787,8 +3873,8 @@ def diff(self, dim, n=1, label='upper'):
         for name, var in self.variables.items():
             if dim in var.dims:
                 if name in self.data_vars:
-                    variables[name] = (var.isel(**kwargs_end) -
-                                       var.isel(**kwargs_start))
+                    variables[name] = (var.isel(**kwargs_end)
+                                       - var.isel(**kwargs_start))
                 else:
                     variables[name] = var.isel(**kwargs_new)
             else:
@@ -3972,8 +4058,8 @@ def sortby(self, variables, ascending=True):
         for data_array in aligned_other_vars:
             if data_array.ndim != 1:
                 raise ValueError("Input DataArray is not 1-D.")
-            if (data_array.dtype == object and
-                    LooseVersion(np.__version__) < LooseVersion('1.11.0')):
+            if (data_array.dtype == object
+                    and LooseVersion(np.__version__) < LooseVersion('1.11.0')):
                 raise NotImplementedError(
                     'sortby uses np.lexsort under the hood, which requires '
                     'numpy 1.11.0 or later to support object data-type.')
@@ -4049,9 +4135,9 @@ def quantile(self, q, dim=None, interpolation='linear',
             reduce_dims = [d for d in var.dims if d in dims]
             if reduce_dims or not var.dims:
                 if name not in self.coords:
-                    if (not numeric_only or
-                        np.issubdtype(var.dtype, np.number) or
-                            var.dtype == np.bool_):
+                    if (not numeric_only
+                        or np.issubdtype(var.dtype, np.number)
+                            or var.dtype == np.bool_):
                         if len(reduce_dims) == var.ndim:
                             # prefer to aggregate over axis=None rather than
                             # axis=(0, 1) if they will be equivalent, because
@@ -4167,8 +4253,8 @@ def differentiate(self, coord, edge_order=1, datetime_unit=None):
 
         variables = OrderedDict()
         for k, v in self.variables.items():
-            if (k in self.data_vars and dim in v.dims and
-                    k not in self.coords):
+            if (k in self.data_vars and dim in v.dims
+                    and k not in self.coords):
                 if _contains_datetime_like_objects(v):
                     v = v._to_numeric(datetime_unit=datetime_unit)
                 grad = duck_array_ops.gradient(
@@ -4344,8 +4430,8 @@ def filter_by_attrs(self, **kwargs):
             has_value_flag = False
             for attr_name, pattern in kwargs.items():
                 attr_value = variable.attrs.get(attr_name)
-                if ((callable(pattern) and pattern(attr_value)) or
-                        attr_value == pattern):
+                if ((callable(pattern) and pattern(attr_value))
+                        or attr_value == pattern):
                     has_value_flag = True
                 else:
                     has_value_flag = False

xarray/tests/test_dataset.py

@@ -2111,6 +2111,28 @@ def test_rename_inplace(self):
         # check virtual variables
         assert_array_equal(data['t.dayofyear'], [1, 2, 3])
 
+    def test_rename_dims(self):
+        original = Dataset(
+            {'x': ('x', [0, 1, 2]), 'y': ('x', [10, 11, 12]), 'z': 42})
+        expected = Dataset(
+            {'x': ('x_new', [0, 1, 2]), 'y': ('x_new', [10, 11, 12]), 'z': 42})
+        name_dict = {'x': 'x_new'}
+        actual = original.rename_dims(name_dict)
+        assert_identical(expected, actual)
+        actual_2 = original.rename_dims(**name_dict)
+        assert_identical(expected, actual_2)
+
+    def test_rename_vars(self):
+        original = Dataset(
+            {'x': ('x', [0, 1, 2]), 'y': ('x', [10, 11, 12]), 'z': 42})
+        expected = Dataset(
+            {'x_new': ('x', [0, 1, 2]), 'y': ('x', [10, 11, 12]), 'z': 42})
+        name_dict = {'x': 'x_new'}
+        actual = original.rename_vars(name_dict)
+        assert_identical(expected, actual)
+        actual_2 = original.rename_vars(**name_dict)
+        assert_identical(expected, actual_2)
+
     def test_swap_dims(self):
         original = Dataset({'x': [1, 2, 3], 'y': ('x', list('abc')), 'z': 42})
         expected = Dataset({'z': 42},