diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 16b0cbf4ea1..0f2bf423449 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -46,7 +46,10 @@ Breaking changes New Features ~~~~~~~~~~~~ - Performance improvement when constructing DataArrays. Significantly speeds up repr for Datasets with large number of variables. - By `Deepak Cherian `_ + By `Deepak Cherian `_. +- :py:meth:`DataArray.swap_dims` & :py:meth:`Dataset.swap_dims` now accept dims + in the form of kwargs as well as a dict, like most similar methods. + By `Maximilian Roos `_. Bug fixes ~~~~~~~~~ diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index f062b70aac1..2fef3edbc43 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -1699,7 +1699,9 @@ def rename( new_name_or_name_dict = cast(Hashable, new_name_or_name_dict) return self._replace(name=new_name_or_name_dict) - def swap_dims(self, dims_dict: Mapping[Hashable, Hashable]) -> "DataArray": + def swap_dims( + self, dims_dict: Mapping[Hashable, Hashable] = None, **dims_kwargs + ) -> "DataArray": """Returns a new DataArray with swapped dimensions. Parameters @@ -1708,6 +1710,10 @@ def swap_dims(self, dims_dict: Mapping[Hashable, Hashable]) -> "DataArray": Dictionary whose keys are current dimension names and whose values are new names. + **dim_kwargs : {dim: , ...}, optional + The keyword arguments form of ``dims_dict``. + One of dims_dict or dims_kwargs must be provided. + Returns ------- swapped : DataArray @@ -1749,6 +1755,7 @@ def swap_dims(self, dims_dict: Mapping[Hashable, Hashable]) -> "DataArray": DataArray.rename Dataset.swap_dims """ + dims_dict = either_dict_or_kwargs(dims_dict, dims_kwargs, "swap_dims") ds = self._to_temp_dataset().swap_dims(dims_dict) return self._from_temp_dataset(ds) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 874e26ff465..f8718377104 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -3155,7 +3155,9 @@ def rename_vars( ) return self._replace(variables, coord_names, dims=dims, indexes=indexes) - def swap_dims(self, dims_dict: Mapping[Hashable, Hashable]) -> "Dataset": + def swap_dims( + self, dims_dict: Mapping[Hashable, Hashable] = None, **dims_kwargs + ) -> "Dataset": """Returns a new object with swapped dimensions. Parameters @@ -3164,6 +3166,10 @@ def swap_dims(self, dims_dict: Mapping[Hashable, Hashable]) -> "Dataset": Dictionary whose keys are current dimension names and whose values are new names. + **dim_kwargs : {existing_dim: new_dim, ...}, optional + The keyword arguments form of ``dims_dict``. + One of dims_dict or dims_kwargs must be provided. + Returns ------- swapped : Dataset @@ -3214,6 +3220,8 @@ def swap_dims(self, dims_dict: Mapping[Hashable, Hashable]) -> "Dataset": """ # TODO: deprecate this method in favor of a (less confusing) # rename_dims() method that only renames dimensions. + + dims_dict = either_dict_or_kwargs(dims_dict, dims_kwargs, "swap_dims") for k, v in dims_dict.items(): if k not in self.dims: raise ValueError( diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index adf282ff34c..fc84687511e 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -1639,6 +1639,16 @@ def test_swap_dims(self): expected.indexes[dim_name], actual.indexes[dim_name] ) + # as kwargs + array = DataArray(np.random.randn(3), {"x": list("abc")}, "x") + expected = DataArray(array.values, {"x": ("y", list("abc"))}, dims="y") + actual = array.swap_dims(x="y") + assert_identical(expected, actual) + for dim_name in set().union(expected.indexes.keys(), actual.indexes.keys()): + pd.testing.assert_index_equal( + expected.indexes[dim_name], actual.indexes[dim_name] + ) + # multiindex case idx = pd.MultiIndex.from_arrays([list("aab"), list("yzz")], names=["y1", "y2"]) array = DataArray(np.random.randn(3), {"y": ("x", idx)}, "x") diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index f71b8ec7741..fed9098701b 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -2748,6 +2748,13 @@ def test_swap_dims(self): actual = original.swap_dims({"x": "u"}) assert_identical(expected, actual) + # as kwargs + expected = Dataset( + {"y": ("u", list("abc")), "z": 42}, coords={"x": ("u", [1, 2, 3])} + ) + actual = original.swap_dims(x="u") + assert_identical(expected, actual) + # handle multiindex case idx = pd.MultiIndex.from_arrays([list("aab"), list("yzz")], names=["y1", "y2"]) original = Dataset({"x": [1, 2, 3], "y": ("x", idx), "z": 42})