Skip to content

Commit

Permalink
Update DataArray.rename + docu (#6665)
Browse files Browse the repository at this point in the history
* update dataarray rename incl docu

* update whats-new

* add support for changing name and dims/coords at the same time

* fix runtime typing issue

* Revert "add support for changing name and dims/coords at the same time"

This reverts commit 31d8521.

* enable rename to None again

* fix a typing problem
  • Loading branch information
headtr1ck authored Jul 18, 2022
1 parent 8f983f1 commit 392a614
Show file tree
Hide file tree
Showing 5 changed files with 103 additions and 25 deletions.
3 changes: 3 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,9 @@ Documentation
sizes. In particular, correct the syntax and replace lists with tuples in the
examples. (:issue:`6333`, :pull:`6334`)
By `Stan West <https://github.com/stanwest>`_.
- Mention that ``xr.DataArray.rename`` can rename coordinates.
(:issue:`5458`, :pull:`6665`)
By `Michael Niklas <https://github.com/headtr1ck>`_.

- Added examples to :py:meth:`Dataset.thin` and :py:meth:`DataArray.thin`
By `Emma Marshall <https://github.com/e-marshall>`_.
Expand Down
27 changes: 16 additions & 11 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -1996,17 +1996,17 @@ def rename(
new_name_or_name_dict: Hashable | Mapping[Any, Hashable] | None = None,
**names: Hashable,
) -> DataArray:
"""Returns a new DataArray with renamed coordinates or a new name.
"""Returns a new DataArray with renamed coordinates, dimensions or a new name.
Parameters
----------
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. Otherwise, use the argument
as the new name for this array.
names to new names for coordinates or dimensions. 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.
new names for coordinates or dimensions.
One of new_name_or_name_dict or names must be provided.
Returns
Expand All @@ -2019,16 +2019,21 @@ def rename(
Dataset.rename
DataArray.swap_dims
"""
if names or utils.is_dict_like(new_name_or_name_dict):
new_name_or_name_dict = cast(
Mapping[Hashable, Hashable], new_name_or_name_dict
)
if new_name_or_name_dict is None and not names:
# change name to None?
return self._replace(name=None)
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")
dataset = self._to_temp_dataset().rename(name_dict)
return self._from_temp_dataset(dataset)
else:
new_name_or_name_dict = cast(Hashable, new_name_or_name_dict)
return self._replace(name=new_name_or_name_dict)
if utils.hashable(new_name_or_name_dict) and names:
# change name + dims/coords
dataset = self._to_temp_dataset().rename(names)
dataarray = self._from_temp_dataset(dataset)
return dataarray._replace(name=new_name_or_name_dict)
# only change name
return self._replace(name=new_name_or_name_dict)

def swap_dims(
self: T_DataArray,
Expand Down
6 changes: 3 additions & 3 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -3558,12 +3558,12 @@ def rename(
name_dict: Mapping[Any, Hashable] | None = None,
**names: Hashable,
) -> T_Dataset:
"""Returns a new object with renamed variables and dimensions.
"""Returns a new object with renamed variables, coordinates and dimensions.
Parameters
----------
name_dict : dict-like, optional
Dictionary whose keys are current variable or dimension names and
Dictionary whose keys are current variable, coordinate or dimension names and
whose values are the desired names.
**names : optional
Keyword form of ``name_dict``.
Expand All @@ -3572,7 +3572,7 @@ def rename(
Returns
-------
renamed : Dataset
Dataset with renamed variables and dimensions.
Dataset with renamed variables, coordinates and dimensions.
See Also
--------
Expand Down
2 changes: 1 addition & 1 deletion xarray/core/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@
InterpOptions = Union[Interp1dOptions, InterpolantOptions]

DatetimeUnitOptions = Literal[
"Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns", "ps", "fs", "as", None
"Y", "M", "W", "D", "h", "m", "s", "ms", "us", "μs", "ns", "ps", "fs", "as", None
]

QueryEngineOptions = Literal["python", "numexpr", None]
Expand Down
90 changes: 80 additions & 10 deletions xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -1653,17 +1653,87 @@ def test_reindex_str_dtype(self, dtype) -> None:
assert actual.dtype == expected.dtype

def test_rename(self) -> None:
renamed = self.dv.rename("bar")
assert_identical(renamed.to_dataset(), self.ds.rename({"foo": "bar"}))
assert renamed.name == "bar"

renamed = self.dv.x.rename({"x": "z"}).rename("z")
assert_identical(renamed, self.ds.rename({"x": "z"}).z)
assert renamed.name == "z"
assert renamed.dims == ("z",)

renamed_kwargs = self.dv.x.rename(x="z").rename("z")
assert_identical(renamed, renamed_kwargs)
da = xr.DataArray(
[1, 2, 3], dims="dim", name="name", coords={"coord": ("dim", [5, 6, 7])}
)

# change name
renamed_name = da.rename("name_new")
assert renamed_name.name == "name_new"
expected_name = da.copy()
expected_name.name = "name_new"
assert_identical(renamed_name, expected_name)

# change name to None?
renamed_noname = da.rename(None)
assert renamed_noname.name is None
expected_noname = da.copy()
expected_noname.name = None
assert_identical(renamed_noname, expected_noname)
renamed_noname = da.rename()
assert renamed_noname.name is None
assert_identical(renamed_noname, expected_noname)

# change dim
renamed_dim = da.rename({"dim": "dim_new"})
assert renamed_dim.dims == ("dim_new",)
expected_dim = xr.DataArray(
[1, 2, 3],
dims="dim_new",
name="name",
coords={"coord": ("dim_new", [5, 6, 7])},
)
assert_identical(renamed_dim, expected_dim)

# change dim with kwargs
renamed_dimkw = da.rename(dim="dim_new")
assert renamed_dimkw.dims == ("dim_new",)
assert_identical(renamed_dimkw, expected_dim)

# change coords
renamed_coord = da.rename({"coord": "coord_new"})
assert "coord_new" in renamed_coord.coords
expected_coord = xr.DataArray(
[1, 2, 3], dims="dim", name="name", coords={"coord_new": ("dim", [5, 6, 7])}
)
assert_identical(renamed_coord, expected_coord)

# change coords with kwargs
renamed_coordkw = da.rename(coord="coord_new")
assert "coord_new" in renamed_coordkw.coords
assert_identical(renamed_coordkw, expected_coord)

# change coord and dim
renamed_both = da.rename({"dim": "dim_new", "coord": "coord_new"})
assert renamed_both.dims == ("dim_new",)
assert "coord_new" in renamed_both.coords
expected_both = xr.DataArray(
[1, 2, 3],
dims="dim_new",
name="name",
coords={"coord_new": ("dim_new", [5, 6, 7])},
)
assert_identical(renamed_both, expected_both)

# change coord and dim with kwargs
renamed_bothkw = da.rename(dim="dim_new", coord="coord_new")
assert renamed_bothkw.dims == ("dim_new",)
assert "coord_new" in renamed_bothkw.coords
assert_identical(renamed_bothkw, expected_both)

# change all
renamed_all = da.rename("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
expected_all = xr.DataArray(
[1, 2, 3],
dims="dim_new",
name="name_new",
coords={"coord_new": ("dim_new", [5, 6, 7])},
)
assert_identical(renamed_all, expected_all)

def test_init_value(self) -> None:
expected = DataArray(
Expand Down

0 comments on commit 392a614

Please sign in to comment.