Skip to content

Commit

Permalink
add support for changing name and dims/coords at the same time
Browse files Browse the repository at this point in the history
  • Loading branch information
headtr1ck committed Jun 5, 2022
1 parent 9bd940e commit 31d8521
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 6 deletions.
15 changes: 10 additions & 5 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -1975,12 +1975,11 @@ def rename(
----------
new_name_or_name_dict : str or dict-like, optional
If the argument is dict-like, it used as a mapping from old
names to new names for coordinates or dimensions. Otherwise,
use the argument as the new name for this array.
names to new names. Otherwise, use the argument as the new name
for this array.
**names : Hashable, optional
The keyword arguments form of a mapping from old names to
new names for coordinates or dimensions.
One of new_name_or_name_dict or names must be provided.
new names. One of new_name_or_name_dict or names must be provided.
Returns
-------
Expand All @@ -1995,8 +1994,14 @@ def rename(
if utils.is_dict_like(new_name_or_name_dict) or new_name_or_name_dict is None:
# change dims/coords
name_dict = either_dict_or_kwargs(new_name_or_name_dict, names, "rename")
if self.name in name_dict:
new_name = name_dict[self.name]
name_dict = {k: v for k, v in name_dict.items() if k != self.name}
else:
new_name = self.name
dataset = self._to_temp_dataset().rename(name_dict)
return self._from_temp_dataset(dataset)
dataarray = self._from_temp_dataset(dataset)
return dataarray._replace(name=new_name)
elif utils.hashable(new_name_or_name_dict):
if names:
# change name + dims/coords
Expand Down
11 changes: 10 additions & 1 deletion xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -1699,7 +1699,9 @@ def test_rename(self) -> None:
assert_identical(renamed_bothkw, expected_both)

# change all
renamed_all = da.rename("name_new", dim="dim_new", coord="coord_new")
renamed_all = da.rename(
{"name": "name_new", "dim": "dim_new", "coord": "coord_new"}
)
assert renamed_all.name == "name_new"
assert renamed_all.dims == ("dim_new",)
assert "coord_new" in renamed_all.coords
Expand All @@ -1711,6 +1713,13 @@ def test_rename(self) -> None:
)
assert_identical(renamed_all, expected_all)

# change all with kwargs
renamed_allkw = da.rename("name_new", dim="dim_new", coord="coord_new")
assert renamed_allkw.name == "name_new"
assert renamed_allkw.dims == ("dim_new",)
assert "coord_new" in renamed_allkw.coords
assert_identical(renamed_allkw, expected_all)

def test_init_value(self) -> None:
expected = DataArray(
np.full((3, 4), 3), dims=["x", "y"], coords=[range(3), range(4)]
Expand Down

0 comments on commit 31d8521

Please sign in to comment.