Skip to content

Commit

Permalink
Fix type of .assign_coords (#8495)
Browse files Browse the repository at this point in the history
* Fix type of `.assign_coords`

As discussed in #8455

* Update xarray/core/common.py

Co-authored-by: Benoit Bovy <bbovy@gfz-potsdam.de>

* Generally improve docstring

* .

---------

Co-authored-by: Benoit Bovy <bbovy@gfz-potsdam.de>
  • Loading branch information
max-sixty and benbovy authored Dec 4, 2023
1 parent 50bd8f9 commit 449c31a
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 23 deletions.
34 changes: 16 additions & 18 deletions xarray/core/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -476,7 +476,7 @@ def _calc_assign_results(

def assign_coords(
self,
coords: Mapping[Any, Any] | None = None,
coords: Mapping | None = None,
**coords_kwargs: Any,
) -> Self:
"""Assign new coordinates to this object.
Expand All @@ -486,15 +486,21 @@ def assign_coords(
Parameters
----------
coords : dict-like or None, optional
A dict where the keys are the names of the coordinates
with the new values to assign. If the values are callable, they are
computed on this object and assigned to new coordinate variables.
If the values are not callable, (e.g. a ``DataArray``, scalar, or
array), they are simply assigned. A new coordinate can also be
defined and attached to an existing dimension using a tuple with
the first element the dimension name and the second element the
values for this new coordinate.
coords : mapping of dim to coord, optional
A mapping whose keys are the names of the coordinates and values are the
coordinates to assign. The mapping will generally be a dict or
:class:`Coordinates`.
* If a value is a standard data value — for example, a ``DataArray``,
scalar, or array — the data is simply assigned as a coordinate.
* If a value is callable, it is called with this object as the only
parameter, and the return value is used as new coordinate variables.
* A coordinate can also be defined and attached to an existing dimension
using a tuple with the first element the dimension name and the second
element the values for this new coordinate.
**coords_kwargs : optional
The keyword arguments form of ``coords``.
One of ``coords`` or ``coords_kwargs`` must be provided.
Expand Down Expand Up @@ -595,14 +601,6 @@ def assign_coords(
Attributes:
description: Weather-related data
Notes
-----
Since ``coords_kwargs`` is a dictionary, the order of your arguments
may not be preserved, and so the order of the new variables is not well
defined. Assigning multiple variables within the same ``assign_coords``
is possible, but you cannot reference other variables created within
the same ``assign_coords`` call.
See Also
--------
Dataset.assign
Expand Down
18 changes: 13 additions & 5 deletions xarray/core/coordinates.py
Original file line number Diff line number Diff line change
Expand Up @@ -571,11 +571,18 @@ def assign(self, coords: Mapping | None = None, **coords_kwargs: Any) -> Self:
Parameters
----------
coords : :class:`Coordinates` or mapping of hashable to Any
Mapping from coordinate names to the new values. If a ``Coordinates``
object is passed, its indexes are assigned in the returned object.
Otherwise, a default (pandas) index is created for each dimension
coordinate found in the mapping.
coords : mapping of dim to coord, optional
A mapping whose keys are the names of the coordinates and values are the
coordinates to assign. The mapping will generally be a dict or
:class:`Coordinates`.
* If a value is a standard data value — for example, a ``DataArray``,
scalar, or array — the data is simply assigned as a coordinate.
* A coordinate can also be defined and attached to an existing dimension
using a tuple with the first element the dimension name and the second
element the values for this new coordinate.
**coords_kwargs
The keyword arguments form of ``coords``.
One of ``coords`` or ``coords_kwargs`` must be provided.
Expand Down Expand Up @@ -605,6 +612,7 @@ def assign(self, coords: Mapping | None = None, **coords_kwargs: Any) -> Self:
* y_level_1 (y) int64 0 1 0 1
"""
# TODO: this doesn't support a callable, which is inconsistent with `DataArray.assign_coords`
coords = either_dict_or_kwargs(coords, coords_kwargs, "assign")
new_coords = self.copy()
new_coords.update(coords)
Expand Down

0 comments on commit 449c31a

Please sign in to comment.