new cell measures implementation

Author: malmans2

cf_xarray/accessor.py

@@ -338,27 +338,21 @@ def _get_measure_variable(
 
 def _get_measure(obj: Union[DataArray, Dataset], key: str) -> List[str]:
     """
-    Translate from cell measures ("area" or "volume") to appropriate variable name.
+    Translate from cell measures to appropriate variable name.
     This function interprets the ``cell_measures`` attribute on DataArrays.
 
     Parameters
     ----------
     obj: DataArray, Dataset
         DataArray belonging to the coordinate to be checked
-    key: str, ["area", "volume"]
+    key: str
         key to check for.
 
     Returns
     -------
     List[str], Variable name(s) in parent xarray object that matches axis or coordinate `key`
     """
 
-    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}"
-        )
-
     if isinstance(obj, DataArray):
         obj = obj._to_temp_dataset()
 
@@ -438,7 +432,7 @@ def _build_docstring(func):
     mapper_docstrings = {
         _get_axis_coord: f"One or more of {(_AXIS_NAMES + _COORD_NAMES)!r}",
         _get_axis_coord_single: f"One of {(_AXIS_NAMES + _COORD_NAMES)!r}",
-        _get_measure_variable: f"One of {_CELL_MEASURES!r}",
+        # _get_measure_variable: f"One of {_CELL_MEASURES!r}",
     }
 
     sig = inspect.signature(func)
@@ -653,6 +647,9 @@ class CFAccessor:
 
     def __init__(self, da):
         self._obj = da
+        self._defined_cell_measures = tuple(
+            set(_CELL_MEASURES + tuple(self.cell_measures))
+        )
 
     def _process_signature(
         self,
@@ -833,7 +830,7 @@ def describe(self):
 
         text += "\nCell Measures:\n"
         measures = self.cell_measures
-        for key in _CELL_MEASURES:
+        for key in sorted(self._defined_cell_measures):
             text += f"\t{key}: {measures[key] if key in measures else []}\n"
 
         text += "\nStandard Names:\n"
@@ -930,15 +927,33 @@ def cell_measures(self) -> Dict[str, List[str]]:
         Returns
         -------
         Dictionary of valid cell measure names that can be used with __getitem__ or .cf[key].
-        Will be ("area", "volume") or a subset thereof.
         """
 
-        measures = {
-            key: apply_mapper(_get_measure, self._obj, key, error=False)
-            for key in _CELL_MEASURES
-        }
+        obj = self._obj
+        all_attrs = [
+            da.attrs.get("cell_measures")
+            for da in obj.coords.values()
+            if da.attrs.get("cell_measures")
+        ]
+        if isinstance(obj, DataArray):
+            all_attrs += (
+                [obj.attrs.get("cell_measures")]
+                if obj.attrs.get("cell_measures")
+                else []
+            )
+        elif isinstance(obj, Dataset):
+            all_attrs += [
+                da.attrs.get("cell_measures")
+                for da in obj.data_vars.values()
+                if da.attrs.get("cell_measures")
+            ]
+
+        measures: Dict[str, List[str]] = dict()
+        for attr in all_attrs:
+            for key, value in parse_cell_methods_attr(attr).items():
+                measures[key] = measures.setdefault(key, []) + [value]
 
-        return {k: sorted(v) for k, v in measures.items() if v}
+        return {k: sorted(set(v)) for k, v in measures.items() if v}
 
     def get_standard_names(self) -> List[str]:
 
@@ -1069,7 +1084,7 @@ def check_results(names, k):
                 check_results(names, k)
                 successful[k] = bool(names)
                 coords.extend(names)
-            elif k in _CELL_MEASURES:
+            elif k in self._defined_cell_measures:
                 measure = _get_measure(self._obj, k)
                 check_results(measure, k)
                 successful[k] = bool(measure)

cf_xarray/tests/test_accessor.py

@@ -52,20 +52,28 @@ def test_coordinates():
     assert actual == expected
 
 
-def test_cell_measures():
+def test_cell_measures(capsys):
     ds = airds.copy(deep=True)
     ds["foo"] = xr.DataArray(ds["cell_area"], attrs=dict(standard_name="foo_std_name"))
     ds["air"].attrs["cell_measures"] += " foo_measure: foo"
     assert "foo_std_name" in ds.cf["air_temperature"].cf
 
     ds["air"].attrs["cell_measures"] += " volume: foo"
-    expected = dict(area=["cell_area"], volume=["foo"])
+    expected = dict(area=["cell_area"], foo_measure=["foo"], volume=["foo"])
     actual = ds["air"].cf.cell_measures
     assert actual == expected
 
     actual = ds.cf.cell_measures
     assert actual == expected
 
+    ds.cf.describe()
+    actual = capsys.readouterr().out
+    expected = (
+        "\nCell Measures:\n\tarea: ['cell_area']\n\tfoo_measure: ['foo']\n\tvolume: ['foo']\n"
+        "\nStandard Names:\n\tair_temperature: ['air']\n\tfoo_std_name: ['foo']\n\n"
+    )
+    assert actual.endswith(expected)
+
 
 def test_standard_names():
     expected = dict(

doc/examples/introduction.ipynb

@@ -559,10 +559,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "**Note:** Although it is possible to assign additional coordinates and cell\n",
-    "measures, `.cf.coordinates` and `.cf.cell_measures` only return a subset of\n",
-    "`(\"longitude\", \"latitude\", \"vertical\", \"time\")` and `(\"area\", \"volume\")`,\n",
-    "respectively.\n"
+    "**Note:** Although it is possible to assign additional coordinates,\n",
+    "`.cf.coordinates` and only returns a subset of \n",
+    "`(\"longitude\", \"latitude\", \"vertical\", \"time\")`.\n"
    ]
   },
   {
@@ -927,7 +926,7 @@
    "source": [
     "## Feature: Weight by Cell Measures\n",
     "\n",
-    "`cf_xarray` can weight by cell measure variables `\"area\"` and `\"volume\"` if the\n",
+    "`cf_xarray` can weight by cell measure variables if the\n",
     "appropriate attribute is set\n"
    ]
   },