Fix bug when coordinates is in both attrs and encoding (#95)

* Fix bug when coordinates is in both attrs and encoding

* Cleaner code

* Fixes

Author: dcherian

cf_xarray/accessor.py

@@ -297,15 +297,19 @@ def _get_axis_coord(var: Union[DataArray, Dataset], key: str) -> List[str]:
             f"cf_xarray did not understand key {key!r}. Expected one of {valid_keys!r}"
         )
 
+    search_in = set()
     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)
+        search_in.update(var.encoding["coordinates"].split(" "))
+    if "coordinates" in var.attrs:
+        search_in.update(var.attrs["coordinates"].split(" "))
+    if not search_in:
+        search_in = set(var.coords)
+
+    # maybe only do this for key in _AXIS_NAMES?
+    search_in.update(var.indexes)
 
     results: Set = set()
-    for coord in itertools.chain(search_in, var.indexes):
+    for coord in search_in:
         for criterion, valid_values in coordinate_criteria.items():
             if key in valid_values:
                 expected = valid_values[key]

cf_xarray/tests/datasets.py

@@ -91,3 +91,4 @@
 )
 romsds.coords["z_rho"] = ("s_rho", np.linspace(-100, 0, 10), {"positive": "up"})
 romsds["temp"] = ("s_rho", np.linspace(20, 30, 10), {"coordinates": "z_rho"})
+romsds["temp"].encoding["coordinates"] = "s_rho"

cf_xarray/tests/test_accessor.py

@@ -9,7 +9,7 @@
 import cf_xarray  # noqa
 
 from . import raise_if_dask_computes
-from .datasets import airds, anc, ds_no_attrs, multiple, popds, romsds
+from .datasets import airds, anc, ds_no_attrs, multiple, popds
 
 mpl.use("Agg")
 
@@ -463,5 +463,22 @@ def test_missing_variable_in_coordinates():
 
 
 def test_Z_vs_vertical_ROMS():
+    from .datasets import romsds
+
+    assert_identical(romsds.s_rho.reset_coords(drop=True), romsds.temp.cf["Z"])
+    assert_identical(romsds.z_rho.reset_coords(drop=True), romsds.temp.cf["vertical"])
+
+    romsds = romsds.copy(deep=True)
+
+    romsds.temp.attrs.clear()
+    # look in encoding
+    assert_identical(romsds.s_rho.reset_coords(drop=True), romsds.temp.cf["Z"])
+    with pytest.raises(KeyError):
+        # z_rho is not in .encoding["coordinates"]
+        # so this won't work
+        romsds.temp.cf["vertical"]
+
+    # use .coords if coordinates attribute is not available
+    romsds.temp.encoding.clear()
     assert_identical(romsds.s_rho.reset_coords(drop=True), romsds.temp.cf["Z"])
     assert_identical(romsds.z_rho.reset_coords(drop=True), romsds.temp.cf["vertical"])