Make _guess_coord_axis work with pint arrays (#247)

Author: dcherian

cf_xarray/accessor.py

@@ -220,13 +220,13 @@ def _get_custom_criteria(
     return list(results)
 
 
-def _get_axis_coord(var: Union[DataArray, Dataset], key: str) -> List[str]:
+def _get_axis_coord(obj: Union[DataArray, Dataset], key: str) -> List[str]:
     """
     Translate from axis or coord name to variable name
 
     Parameters
     ----------
-    var : DataArray, Dataset
+    obj : DataArray, Dataset
         DataArray belonging to the coordinate to be checked
     key : str, ["X", "Y", "Z", "T", "longitude", "latitude", "vertical", "time"]
         key to check for.
@@ -260,25 +260,29 @@ def _get_axis_coord(var: Union[DataArray, Dataset], key: str) -> List[str]:
         )
 
     search_in = set()
-    if "coordinates" in var.encoding:
-        search_in.update(var.encoding["coordinates"].split(" "))
-    if "coordinates" in var.attrs:
-        search_in.update(var.attrs["coordinates"].split(" "))
+    if "coordinates" in obj.encoding:
+        search_in.update(obj.encoding["coordinates"].split(" "))
+    if "coordinates" in obj.attrs:
+        search_in.update(obj.attrs["coordinates"].split(" "))
     if not search_in:
-        search_in = set(var.coords)
+        search_in = set(obj.coords)
 
     # maybe only do this for key in _AXIS_NAMES?
-    search_in.update(var.indexes)
+    search_in.update(obj.indexes)
 
+    search_in = search_in & set(obj.coords)
     results: Set = set()
     for coord in search_in:
+        var = obj.coords[coord]
         if key in coordinate_criteria:
             for criterion, expected in coordinate_criteria[key].items():
-                if (
-                    coord in var.coords
-                    and var.coords[coord].attrs.get(criterion, None) in expected
-                ):
+                if var.attrs.get(criterion, None) in expected:
                     results.update((coord,))
+                if criterion == "units":
+                    # deal with pint-backed objects
+                    units = getattr(var.data, "units", None)
+                    if units in expected:
+                        results.update((coord,))
     return list(results)
 
 

cf_xarray/tests/__init__.py

@@ -1,5 +1,7 @@
+import importlib
 import re
 from contextlib import contextmanager
+from distutils import version
 
 import dask
 import pytest
@@ -39,3 +41,26 @@ def __call__(self, dsk, keys, **kwargs):
 def raise_if_dask_computes(max_computes=0):
     scheduler = CountingScheduler(max_computes)
     return dask.config.set(scheduler=scheduler)
+
+
+def _importorskip(modname, minversion=None):
+    try:
+        mod = importlib.import_module(modname)
+        has = True
+        if minversion is not None:
+            if LooseVersion(mod.__version__) < LooseVersion(minversion):
+                raise ImportError("Minimum version not satisfied")
+    except ImportError:
+        has = False
+    func = pytest.mark.skipif(not has, reason=f"requires {modname}")
+    return has, func
+
+
+def LooseVersion(vstring):
+    # Our development version is something like '0.10.9+aac7bfc'
+    # This function just ignored the git commit id.
+    vstring = vstring.split("+")[0]
+    return version.LooseVersion(vstring)
+
+
+has_pint, requires_pint = _importorskip("pint")

cf_xarray/tests/test_accessor.py

@@ -24,7 +24,7 @@
     popds,
     romsds,
 )
-from . import raise_if_dask_computes
+from . import raise_if_dask_computes, requires_pint
 
 mpl.use("Agg")
 
@@ -149,6 +149,19 @@ def test_coordinates():
     assert actual == expected
 
 
+@requires_pint
+def test_coordinates_quantified():
+    # note: import order is important
+    from .. import units  # noqa
+
+    pytest.importorskip("pint_xarray")
+
+    quantified = popds.pint.quantify()
+    assert_identical(
+        quantified.cf[["latitude"]].pint.dequantify(), popds.cf[["latitude"]]
+    )
+
+
 def test_cell_measures():
     ds = airds.copy(deep=True)
     ds["foo"] = xr.DataArray(ds["cell_area"], attrs=dict(standard_name="foo_std_name"))

doc/whats-new.rst

@@ -3,8 +3,9 @@
 What's New
 ----------
 
-v0.6.0 (unreleased)
+v0.6.1 (unreleased)
 ===================
+- Support detecting pint-backed Variables with units-based criteria. By `Deepak Cherian`_.
 
 v0.6.0 (June 29, 2021)
 ======================