add_bounds uses keys rather than dims (#221)

Author: malmans2

cf_xarray/accessor.py

@@ -1586,15 +1586,15 @@ def get_bounds_dim_name(self, key: str) -> str:
         assert self._obj.sizes[bounds_dim] in [2, 4]
         return bounds_dim
 
-    def add_bounds(self, dims: Union[Hashable, Iterable[Hashable]]):
+    def add_bounds(self, keys: Union[str, Iterable[str]]):
         """
         Returns a new object with bounds variables. The bounds values are guessed assuming
         equal spacing on either side of a coordinate label.
 
         Parameters
         ----------
-        dims : Hashable or Iterable[Hashable]
-            Either a single dimension name or a list of dimension names.
+        keys : str or Iterable[str]
+            Either a single key or a list of keys corresponding to dimensions.
 
         Returns
         -------
@@ -1609,12 +1609,16 @@ def add_bounds(self, dims: Union[Hashable, Iterable[Hashable]]):
         The bounds variables are automatically named f"{dim}_bounds" where ``dim``
         is a dimension name.
         """
-        if isinstance(dims, Hashable):
-            dimensions = (dims,)
-        else:
-            dimensions = dims
+        if isinstance(keys, str):
+            keys = [keys]
+
+        dimensions = set()
+        for key in keys:
+            dimensions.update(
+                apply_mapper(_get_dims, self._obj, key, error=False, default=[key])
+            )
 
-        bad_dims: Set[Hashable] = set(dimensions) - set(self._obj.dims)
+        bad_dims: Set[str] = dimensions - set(self._obj.dims)
         if bad_dims:
             raise ValueError(
                 f"{bad_dims!r} are not dimensions in the underlying object."

cf_xarray/tests/test_accessor.py

@@ -634,6 +634,10 @@ def test_add_bounds(obj, dims):
         assert added[dim].attrs["bounds"] == name
         assert_allclose(added[name].reset_coords(drop=True), expected[dim])
 
+    # Test multiple dimensions
+    assert not {"x1_bounds", "x2_bounds"} <= set(multiple.variables)
+    assert {"x1_bounds", "x2_bounds"} <= set(multiple.cf.add_bounds("X").variables)
+
 
 def test_bounds():
     ds = airds.copy(deep=True).cf.add_bounds("lat")

doc/whats-new.rst

@@ -6,6 +6,7 @@ What's New
 v0.5.2 (unreleased)
 ===================
 
+- Replace the ``dims`` argument of :py:meth:`Dataset.cf.add_bounds` with ``keys``, allowing to use CF keys. By `Mattia Almansi`_.
 - Added :py:attr:`DataArray.cf.formula_terms` and :py:attr:`Dataset.cf.formula_terms`.
   By `Deepak Cherian`_.
 - Added :py:attr:`Dataset.cf.bounds` to return a dictionary mapping valid keys to the variable names of their bounds. By `Mattia Almansi`_.