From a22251b8b93e396a7dd8e8739be71797d4173ad6 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Thu, 30 Nov 2023 13:57:29 -0800 Subject: [PATCH 1/6] Fix type of `.assign_coords` As discussed in #8455 --- xarray/core/common.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index cb5b79defc0..a64c198733f 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -11,6 +11,7 @@ import pandas as pd from xarray.core import dtypes, duck_array_ops, formatting, formatting_html, ops +from xarray.core.coordinates import Coordinates from xarray.core.indexing import BasicIndexer, ExplicitlyIndexed from xarray.core.options import OPTIONS, _get_keep_attrs from xarray.core.parallelcompat import get_chunked_array_type, guess_chunkmanager @@ -474,7 +475,7 @@ def _calc_assign_results( def assign_coords( self, - coords: Mapping[Any, Any] | None = None, + coords: Mapping[Any, Any] | Coordinates | None = None, **coords_kwargs: Any, ) -> Self: """Assign new coordinates to this object. @@ -484,7 +485,7 @@ def assign_coords( Parameters ---------- - coords : dict-like or None, optional + coords : dict-like or Coordinates, 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. @@ -593,14 +594,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 From 20267b66b4772482187962c1b8cc451840a6a5c5 Mon Sep 17 00:00:00 2001 From: Maximilian Roos <5635139+max-sixty@users.noreply.github.com> Date: Fri, 1 Dec 2023 12:28:12 -0800 Subject: [PATCH 2/6] Update xarray/core/common.py Co-authored-by: Benoit Bovy --- xarray/core/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index a64c198733f..592ed66e936 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -475,7 +475,7 @@ def _calc_assign_results( def assign_coords( self, - coords: Mapping[Any, Any] | Coordinates | None = None, + coords: Mapping | None = None, **coords_kwargs: Any, ) -> Self: """Assign new coordinates to this object. From efc6ee5bb758f92a2537aa0507d9986b3947f549 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Fri, 1 Dec 2023 13:14:32 -0800 Subject: [PATCH 3/6] Generally improve docstring --- xarray/core/common.py | 20 +++++++++++--------- xarray/core/coordinates.py | 15 ++++++++++----- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index 592ed66e936..8321e187f15 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -485,15 +485,17 @@ def assign_coords( Parameters ---------- - coords : dict-like or Coordinates, 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 use 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. diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index 0c85b2a2d69..a13dd549332 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -571,11 +571,15 @@ 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 +609,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) From f8aff096ccd1e8aa026809ae62805bf88ddfb588 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Fri, 1 Dec 2023 13:15:31 -0800 Subject: [PATCH 4/6] . --- xarray/core/common.py | 1 - 1 file changed, 1 deletion(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index 8321e187f15..7d75650c13d 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -11,7 +11,6 @@ import pandas as pd from xarray.core import dtypes, duck_array_ops, formatting, formatting_html, ops -from xarray.core.coordinates import Coordinates from xarray.core.indexing import BasicIndexer, ExplicitlyIndexed from xarray.core.options import OPTIONS, _get_keep_attrs from xarray.core.parallelcompat import get_chunked_array_type, guess_chunkmanager From 5dde850dfda2e54b0751e91458dd4ce0b7e49ffb Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Fri, 1 Dec 2023 21:47:09 -0800 Subject: [PATCH 5/6] --- xarray/core/common.py | 8 ++++---- xarray/core/coordinates.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index 7d75650c13d..36aae1c6aa7 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -488,11 +488,11 @@ def assign_coords( 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``, + * 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 use as new coordinate variables. - - A coordinate can also be defined and attached to an existing dimension + * 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 diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index a13dd549332..bdd94905efb 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -575,9 +575,9 @@ def assign(self, coords: Mapping | None = None, **coords_kwargs: Any) -> Self: 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``, + * 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 + * 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 From 83c7a7a4ca6a0f18bb7c977cc58e2577ae47c08b Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sat, 2 Dec 2023 11:20:03 -0800 Subject: [PATCH 6/6] --- xarray/core/common.py | 4 ++++ xarray/core/coordinates.py | 3 +++ 2 files changed, 7 insertions(+) diff --git a/xarray/core/common.py b/xarray/core/common.py index 36aae1c6aa7..3a9b1bf2d48 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -488,13 +488,17 @@ def assign_coords( 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. diff --git a/xarray/core/coordinates.py b/xarray/core/coordinates.py index bdd94905efb..cdf1d354be6 100644 --- a/xarray/core/coordinates.py +++ b/xarray/core/coordinates.py @@ -575,11 +575,14 @@ def assign(self, coords: Mapping | None = None, **coords_kwargs: Any) -> Self: 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.