Fix type of .assign_coords (pydata#8495)

* Fix type of `.assign_coords` As discussed in pydata#8455 * Update xarray/core/common.py Co-authored-by: Benoit Bovy <[email protected]> * Generally improve docstring * . --------- Co-authored-by: Benoit Bovy <[email protected]>
headtr1ck · Dec 4, 2023 · 449c31a · 449c31a
1 parent 50bd8f9
commit 449c31a
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 23 deletions.
diff --git 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
@@ -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)