Look for & prioritize "coordinates" attribute when searching. (#16)

* Look for & prioritize "coordinates" attribute when searching.

* fix tests.

* Apply suggestions from code review

Author: dcherian

cf_xarray/accessor.py

@@ -1,6 +1,6 @@
 import functools
 import inspect
-from typing import Any, Union
+from typing import Any, List, Optional, Set, Union
 
 import xarray as xr
 from xarray import DataArray, Dataset
@@ -44,6 +44,7 @@
         "Y": ("latitude",),
         "X": ("longitude",),
     },
+    "long_name": {"T": ("time",)},
     "_CoordinateAxisType": {
         "T": ("Time",),
         "Z": ("GeoZ", "Height", "Pressure"),
@@ -84,13 +85,29 @@
 }
 
 
-def _get_axis_coord(var: xr.DataArray, key, error: bool = True, default: Any = None):
+def _get_axis_coord_single(var, key, *args):
+    """ Helper method for when we really want only one result per key. """
+    results = _get_axis_coord(var, key, *args)
+    if len(results) > 1:
+        raise ValueError(
+            "Multiple results for {key!r} found: {results!r}. Is this valid CF? Please open an issue."
+        )
+    else:
+        return results[0]
+
+
+def _get_axis_coord(
+    var: Union[xr.DataArray, xr.Dataset],
+    key: str,
+    error: bool = True,
+    default: Optional[str] = None,
+) -> List[Optional[str]]:
     """
     Translate from axis or coord name to variable name
 
     Parameters
     ----------
-    var : `xarray.DataArray`
+    var : DataArray, Dataset
         DataArray belonging to the coordinate to be checked
     key : str, ["X", "Y", "Z", "T", "longitude", "latitude", "vertical", "time"]
         key to check for.
@@ -102,7 +119,7 @@ def _get_axis_coord(var: xr.DataArray, key, error: bool = True, default: Any = N
 
     Returns
     -------
-    str, Variable name in parent xarray object that matches axis or coordinate `key`
+    List[str], Variable name(s) in parent xarray object that matches axis or coordinate `key`
 
     Notes
     -----
@@ -128,22 +145,33 @@ def _get_axis_coord(var: xr.DataArray, key, error: bool = True, default: Any = N
         if error:
             raise KeyError(f"Did not understand {key}")
         else:
-            return default
+            return [default]
 
     if axis is None:
         raise AssertionError("Should be unreachable")
 
-    for coord in var.coords:
+    if "coordinates" in var.encoding:
+        search_in = var.encoding["coordinates"].split(" ")
+    elif "coordinates" in var.attrs:
+        search_in = var.attrs["coordinates"].split(" ")
+    else:
+        search_in = list(var.coords)
+
+    results: Set = set()
+    for coord in search_in:
         for criterion, valid_values in coordinate_criteria.items():
             if axis in valid_values:  # type: ignore
                 expected = valid_values[axis]  # type: ignore
                 if var.coords[coord].attrs.get(criterion, None) in expected:
-                    return coord
+                    results.update((coord,))
 
-    if error:
-        raise KeyError(f"axis name {key!r} not found!")
+    if not results:
+        if error:
+            raise KeyError(f"axis name {key!r} not found!")
+        else:
+            return [default]
     else:
-        return default
+        return list(results)
 
 
 def _get_measure_variable(
@@ -184,7 +212,9 @@ def _get_measure(da: xr.DataArray, key: str, error: bool = True, default: Any =
     return measures[key]
 
 
-_DEFAULT_KEY_MAPPERS: dict = dict.fromkeys(("dim", "coord", "group"), _get_axis_coord)
+_DEFAULT_KEY_MAPPERS: dict = dict.fromkeys(
+    ("dim", "coord", "group"), _get_axis_coord_single
+)
 _DEFAULT_KEY_MAPPERS["weights"] = _get_measure_variable
 
 
@@ -261,7 +291,7 @@ def __call__(self, *args, **kwargs):
             obj=self._obj,
             attr="plot",
             accessor=self.accessor,
-            key_mappers=dict.fromkeys(self._keys, _get_axis_coord),
+            key_mappers=dict.fromkeys(self._keys, _get_axis_coord_single),
         )
         return plot(*args, **kwargs)
 
@@ -270,7 +300,7 @@ def __getattr__(self, attr):
             obj=self._obj.plot,
             attr=attr,
             accessor=self.accessor,
-            key_mappers=dict.fromkeys(self._keys, _get_axis_coord),
+            key_mappers=dict.fromkeys(self._keys, _get_axis_coord_single),
         )
 
 
@@ -294,7 +324,6 @@ def _process_signature(self, func, args, kwargs, key_mappers):
             arguments = self._rewrite_values(
                 bound.arguments, key_mappers, tuple(var_kws)
             )
-            print(arguments)
         else:
             arguments = {}
 
@@ -311,7 +340,7 @@ def _process_signature(self, func, args, kwargs, key_mappers):
     def _rewrite_values(self, kwargs, key_mappers: dict, var_kws):
         """ rewrites 'dim' for example using 'mapper' """
         updates: dict = {}
-        key_mappers.update(dict.fromkeys(var_kws, _get_axis_coord))
+        key_mappers.update(dict.fromkeys(var_kws, _get_axis_coord_single))
         for key, mapper in key_mappers.items():
             value = kwargs.get(key, None)
             if value is not None:
@@ -341,7 +370,7 @@ def _rewrite_values(self, kwargs, key_mappers: dict, var_kws):
         for vkw in var_kws:
             if vkw in kwargs:
                 maybe_update = {
-                    k: _get_axis_coord(self._obj, v, False, v)
+                    k: _get_axis_coord_single(self._obj, v, False, v)
                     for k, v in kwargs[vkw].items()
                     if k in key_mappers
                 }
@@ -367,22 +396,21 @@ def plot(self):
 class CFDatasetAccessor(CFAccessor):
     def __getitem__(self, key):
         if key in _AXIS_NAMES + _COORD_NAMES:
-            return self._obj[_get_axis_coord(self._obj, key)]
+            varnames = _get_axis_coord(self._obj, key)
+            return self._obj.reset_coords()[varnames].set_coords(varnames)
         elif key in _CELL_MEASURES:
             raise NotImplementedError("measures not implemented for Dataset yet.")
             # return self._obj[_get_measure(self._obj)[key]]
         else:
             raise KeyError(f"DataArray.cf does not understand the key {key}")
 
-    # def __getitem__(self, key):
-    #     raise AttributeError("Dataset.cf does not support [] indexing or __getitem__")
-
 
 @xr.register_dataarray_accessor("cf")
 class CFDataArrayAccessor(CFAccessor):
     def __getitem__(self, key):
         if key in _AXIS_NAMES + _COORD_NAMES:
-            return self._obj[_get_axis_coord(self._obj, key)]
+            varname = _get_axis_coord_single(self._obj, key)
+            return self._obj[varname].reset_coords(drop=True)
         elif key in _CELL_MEASURES:
             return self._obj[_get_measure(self._obj, key)]
         else:

cf_xarray/tests/test_accessor.py

@@ -135,6 +135,8 @@ def test_dataset_plot(obj):
 )
 def test_getitem(obj, key, expected_key):
     actual = obj.cf[key]
+    if isinstance(obj, xr.Dataset):
+        expected_key = [expected_key]
     expected = obj[expected_key]
     assert_identical(actual, expected)
 
@@ -146,3 +148,31 @@ def test_getitem_errors(obj,):
     obj.lon.attrs = {}
     with pytest.raises(KeyError):
         obj.cf["X"]
+
+
+def test_getitem_uses_coordinates():
+    # POP-like dataset
+    ds = xr.Dataset()
+    ds.coords["TLONG"] = (("nlat", "nlon"), np.ones((20, 30)), {"axis": "X"})
+    ds.coords["TLAT"] = (("nlat", "nlon"), 2 * np.ones((20, 30)), {"axis": "Y"})
+    ds.coords["ULONG"] = (("nlat", "nlon"), 0.5 * np.ones((20, 30)), {"axis": "X"})
+    ds.coords["ULAT"] = (("nlat", "nlon"), 2.5 * np.ones((20, 30)), {"axis": "Y"})
+    ds["UVEL"] = (
+        ("nlat", "nlon"),
+        np.ones((20, 30)) * 15,
+        {"coordinates": "ULONG ULAT"},
+    )
+    ds["TEMP"] = (
+        ("nlat", "nlon"),
+        np.ones((20, 30)) * 15,
+        {"coordinates": "TLONG TLAT"},
+    )
+
+    assert_identical(
+        ds.cf["X"], ds.reset_coords()[["ULONG", "TLONG"]].set_coords(["ULONG", "TLONG"])
+    )
+    assert_identical(
+        ds.cf["Y"], ds.reset_coords()[["ULAT", "TLAT"]].set_coords(["ULAT", "TLAT"])
+    )
+    assert_identical(ds.UVEL.cf["X"], ds["ULONG"].reset_coords(drop=True))
+    assert_identical(ds.TEMP.cf["X"], ds["TLONG"].reset_coords(drop=True))