Support indexing by cf_role for DSG datasets (#240)

xarray-contrib · Jun 29, 2021 · 2f46ee3 · 2f46ee3
1 parent ed36d29
commit 2f46ee3
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/cf_xarray/accessor.py b/cf_xarray/accessor.py
@@ -25,7 +25,7 @@
 from xarray import DataArray, Dataset
 from xarray.core.arithmetic import SupportsArithmetic
 
-from .criteria import coordinate_criteria, regex
+from .criteria import cf_role_criteria, coordinate_criteria, regex
 from .helpers import bounds_to_vertices
 from .utils import (
     _get_version,
@@ -372,6 +372,7 @@ def _get_all(obj: Union[DataArray, Dataset], key: str) -> List[str]:
     """
     all_mappers = (
         _get_custom_criteria,
+        functools.partial(_get_custom_criteria, criteria=cf_role_criteria),
         _get_axis_coord,
         _get_measure,
         _get_with_standard_name,
@@ -665,7 +666,7 @@ def check_results(names, key):
             successful[k] = bool(measure)
             if measure:
                 varnames.extend(measure)
-        elif k in custom_criteria:
+        elif k in custom_criteria or k in cf_role_criteria:
             names = _get_all(obj, k)
             check_results(names, k)
             successful[k] = bool(names)

diff --git a/cf_xarray/criteria.py b/cf_xarray/criteria.py
@@ -6,7 +6,11 @@
 
 
 import re
-from typing import MutableMapping, Tuple
+from typing import Mapping, MutableMapping, Tuple
+
+cf_role_criteria: Mapping[str, Mapping[str, str]] = {
+    k: {"cf_role": k} for k in ("timeseries_id", "profile_id", "trajectory_id")
+}
 
 coordinate_criteria: MutableMapping[str, MutableMapping[str, Tuple]] = {
     "latitude": {

diff --git a/doc/whats-new.rst b/doc/whats-new.rst
@@ -5,6 +5,7 @@ What's New
 
 v0.5.3 (unreleased)
 ===================
+- Support indexing by ``cf_role`` attribute. By `Deepak Cherian`_.
 - Implemented :py:meth:`Dataset.cf.add_canonical_attributes` and :py:meth:`DataArray.cf.add_canonical_attributes`
   to add CF canonical attributes. By `Mattia Almansi`_.
 - Begin adding support for units with a unit registry for pint arrays. :pr:`197`.