Support Dataset.cf.cell_measures (#117)

* Support Dataset.cf.cell_measures

* Cleanup

* fix docstring

Author: dcherian

cf_xarray/accessor.py

@@ -336,45 +336,44 @@ def _get_measure_variable(
     return [da[varnames[0]]]
 
 
-def _get_measure(da: Union[DataArray, Dataset], key: str) -> List[str]:
+def _get_measure(obj: Union[DataArray, Dataset], key: str) -> List[str]:
     """
     Translate from cell measures ("area" or "volume") to appropriate variable name.
     This function interprets the ``cell_measures`` attribute on DataArrays.
 
     Parameters
     ----------
-    da: DataArray
+    obj: DataArray, Dataset
         DataArray belonging to the coordinate to be checked
     key: str, ["area", "volume"]
         key to check for.
-    error: bool
-        raise errors when key is not found or interpretable. Use False and provide default
-        to replicate dict.get(k, None).
-    default: Any
-        default value to return when error is False.
 
     Returns
     -------
     List[str], Variable name(s) in parent xarray object that matches axis or coordinate `key`
     """
-    if not isinstance(da, DataArray):
-        raise NotImplementedError("Measures not implemented for Datasets yet.")
-
-    if "cell_measures" not in da.attrs:
-        raise KeyError("'cell_measures' not present in 'attrs'.")
 
     valid_keys = _CELL_MEASURES
     if key not in valid_keys:
         raise KeyError(
             f"cf_xarray did not understand key {key!r}. Expected one of {valid_keys!r}"
         )
 
-    attr = da.attrs["cell_measures"]
-    measures = parse_cell_methods_attr(attr)
-    results: Union[str, List] = measures.get(key, [])
+    if isinstance(obj, DataArray):
+        obj = obj._to_temp_dataset()
+
+    results = set()
+    for var in obj.variables:
+        da = obj[var]
+        if "cell_measures" in da.attrs:
+            attr = da.attrs["cell_measures"]
+            measures = parse_cell_methods_attr(attr)
+            if key in measures:
+                results.update([measures[key]])
+
     if isinstance(results, str):
         return [results]
-    return results
+    return list(results)
 
 
 #: Default mappers for common keys.
@@ -833,10 +832,7 @@ def describe(self):
             text += f"\t{key}: {coords[key] if key in coords else []}\n"
 
         text += "\nCell Measures:\n"
-        if isinstance(self._obj, Dataset):
-            measures = {key: "unsupported" for key in _CELL_MEASURES}
-        else:
-            measures = self.cell_measures
+        measures = self.cell_measures
         for key in _CELL_MEASURES:
             text += f"\t{key}: {measures[key] if key in measures else []}\n"
 
@@ -930,14 +926,12 @@ def cell_measures(self) -> Dict[str, List[str]]:
 
         This is useful for checking whether a key is valid for indexing, i.e.
         that the attributes necessary to allow indexing by that key exist.
-        However, it will only return the cell measure names.
 
         Returns
         -------
         Dictionary of valid cell measure names that can be used with __getitem__ or .cf[key].
         Will be ("area", "volume") or a subset thereof.
         """
-        assert isinstance(self._obj, DataArray), "this only works with DataArrays"
 
         measures = {
             key: apply_mapper(_get_measure, self._obj, key, error=False)

cf_xarray/tests/test_accessor.py

@@ -26,7 +26,7 @@ def test_describe(capsys):
         "Axes:\n\tX: ['lon']\n\tY: ['lat']\n\tZ: []\n\tT: ['time']\n"
         "\nCoordinates:\n\tlongitude: ['lon']\n\tlatitude: ['lat']"
         "\n\tvertical: []\n\ttime: ['time']\n"
-        "\nCell Measures:\n\tarea: unsupported\n\tvolume: unsupported\n"
+        "\nCell Measures:\n\tarea: ['cell_area']\n\tvolume: []\n"
         "\nStandard Names:\n\tair_temperature: ['air']\n\n"
     )
     assert actual == expected
@@ -58,12 +58,13 @@ def test_cell_measures():
     ds["air"].attrs["cell_measures"] += " foo_measure: foo"
     assert "foo_std_name" in ds.cf["air_temperature"].cf
 
-    expected = dict(area=["cell_area"])
+    ds["air"].attrs["cell_measures"] += " volume: foo"
+    expected = dict(area=["cell_area"], volume=["foo"])
     actual = ds["air"].cf.cell_measures
     assert actual == expected
 
-    with pytest.raises(AssertionError, match=r"this only works with DataArrays"):
-        popds.cf.cell_measures
+    actual = ds.cf.cell_measures
+    assert actual == expected
 
 
 def test_standard_names():

doc/whats-new.rst

@@ -4,6 +4,7 @@ What's New
 
 v0.4.0 (unreleased)
 ===================
+- Support ``Dataset.cf.cell_measures``. By `Deepak Cherian`_.
 - Added ``.axes`` to return a dictionary mapping available Axis standard names to variable names of an xarray object, ``.coordinates`` for Coordinates, 
   ``.cell_measures`` for Cell Measures, and ``.standard_names`` for all variables. `Kristen Thyng`_ and `Mattia Almansi`_.
 - Changed ``get_valid_keys()`` to ``.keys()``. `Kristen Thyng`_.