Allow adding bounds to general coordinate vars (#293)

Author: dcherian

cf_xarray/accessor.py

@@ -2057,7 +2057,7 @@ def add_bounds(self, keys: str | Iterable[str]):
         Parameters
         ----------
         keys : str or Iterable[str]
-            Either a single key or a list of keys corresponding to dimensions.
+            Either a single variable name or a list of variable names.
 
         Returns
         -------
@@ -2070,8 +2070,8 @@ def add_bounds(self, keys: str | Iterable[str]):
 
         Notes
         -----
-        The bounds variables are automatically named ``f"{dim}_bounds"`` where ``dim``
-        is a dimension name.
+        The bounds variables are automatically named ``f"{var}_bounds"`` where ``var``
+        is a variable name.
 
         Examples
         --------
@@ -2085,25 +2085,26 @@ def add_bounds(self, keys: str | Iterable[str]):
         if isinstance(keys, str):
             keys = [keys]
 
-        dimensions = set()
+        variables = set()
         for key in keys:
-            dimensions.update(
-                apply_mapper(_get_dims, self._obj, key, error=False, default=[key])
+            variables.update(
+                apply_mapper(_get_all, self._obj, key, error=False, default=[key])
             )
 
-        bad_dims: set[str] = dimensions - set(self._obj.dims)
-        if bad_dims:
+        obj = self._maybe_to_dataset(self._obj.copy(deep=True))
+
+        bad_vars: set[str] = variables - set(obj.variables)
+        if bad_vars:
             raise ValueError(
-                f"{bad_dims!r} are not dimensions in the underlying object."
+                f"{bad_vars!r} are not variables in the underlying object."
             )
 
-        obj = self._maybe_to_dataset(self._obj.copy(deep=True))
-        for dim in dimensions:
-            bname = f"{dim}_bounds"
+        for var in variables:
+            bname = f"{var}_bounds"
             if bname in obj.variables:
                 raise ValueError(f"Bounds variable name {bname!r} will conflict!")
-            obj.coords[bname] = _guess_bounds_dim(obj[dim].reset_coords(drop=True))
-            obj[dim].attrs["bounds"] = bname
+            obj.coords[bname] = _guess_bounds_dim(obj[var].reset_coords(drop=True))
+            obj[var].attrs["bounds"] = bname
 
         return self._maybe_to_dataarray(obj)
 

cf_xarray/tests/test_accessor.py

@@ -699,9 +699,10 @@ def test_plot_xincrease_yincrease():
         assert lim[0] > lim[1]
 
 
-@pytest.mark.parametrize("dims", ["lat", "time", ["lat", "lon"]])
-@pytest.mark.parametrize("obj", [airds])
-def test_add_bounds(obj, dims):
+@pytest.mark.parametrize("dims", ["time2", "lat", "time", ["lat", "lon"]])
+def test_add_bounds(dims):
+    obj = airds.copy(deep=True)
+
     expected = {}
     expected["lat"] = xr.concat(
         [
@@ -728,10 +729,12 @@ def test_add_bounds(obj, dims):
         ],
         dim="bounds",
     )
+    expected["time2"] = expected["time"]
     expected["lat"].attrs.clear()
     expected["lon"].attrs.clear()
     expected["time"].attrs.clear()
 
+    obj.coords["time2"] = obj.time
     added = obj.cf.add_bounds(dims)
     if isinstance(dims, str):
         dims = (dims,)
@@ -742,6 +745,8 @@ def test_add_bounds(obj, dims):
         assert added[dim].attrs["bounds"] == name
         assert_allclose(added[name].reset_coords(drop=True), expected[dim])
 
+
+def test_add_bounds_multiple():
     # Test multiple dimensions
     assert not {"x1_bounds", "x2_bounds"} <= set(multiple.variables)
     assert {"x1_bounds", "x2_bounds"} <= set(multiple.cf.add_bounds("X").variables)