From 449c31a70447740bffa4dc7b471ea0d18da7f7c0 Mon Sep 17 00:00:00 2001 From: Maximilian Roos <5635139+max-sixty@users.noreply.github.com> Date: Mon, 4 Dec 2023 11:11:55 -0800 Subject: [PATCH] Fix type of `.assign_coords` (#8495) * Fix type of `.assign_coords` As discussed in #8455 * Update xarray/core/common.py Co-authored-by: Benoit Bovy * Generally improve docstring * . --------- Co-authored-by: Benoit Bovy --- xarray/core/common.py | 34 ++++++++++++++++------------------ xarray/core/coordinates.py | 18 +++++++++++++----- 2 files changed, 29 insertions(+), 23 deletions(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index cebd8f2a95b..c2c0a79edb6 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -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. @@ -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. @@ -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 diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 0c85b2a2d69..cdf1d354be6 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -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. @@ -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)