Add getitem and actually parse attrs (#12)

* Add getitem and actually parse attrs

Closes #3
Partially implements #4

* lint.

* Better test coverage.

Author: dcherian

cf_xarray/accessor.py

@@ -1,6 +1,6 @@
 import functools
 import inspect
-from typing import Union
+from typing import Any, Union
 
 import xarray as xr
 from xarray import DataArray, Dataset
@@ -15,10 +15,149 @@
 
 
 _DEFAULT_KEYS_TO_REWRITE = ("dim", "coord", "group")
+_AXIS_NAMES = ("X", "Y", "Z", "T")
+_COORD_NAMES = ("longitude", "latitude", "vertical", "time")
+_COORD_AXIS_MAPPING = dict(zip(_COORD_NAMES, _AXIS_NAMES))
+_CELL_MEASURES = ("area", "volume")
+
+
+# Define the criteria for coordinate matches
+# Copied from metpy
+# Internally we only use X, Y, Z, T
+# TODO: Metpy adds latitude and longitude separately so we may revert to doing that too
+coordinate_criteria = {
+    "standard_name": {
+        "T": ("time",),
+        "Z": (
+            "air_pressure",
+            "height",
+            "geopotential_height",
+            "altitude",
+            "model_level_number",
+            "atmosphere_ln_pressure_coordinate",
+            "atmosphere_sigma_coordinate",
+            "atmosphere_hybrid_sigma_pressure_coordinate",
+            "atmosphere_hybrid_height_coordinate",
+            "atmosphere_sleve_coordinate",
+            "height_above_geopotential_datum",
+            "height_above_reference_ellipsoid",
+            "height_above_mean_sea_level",
+        ),
+        "Y": ("latitude",),
+        "X": ("longitude",),
+    },
+    "_CoordinateAxisType": {
+        "T": ("Time",),
+        "Z": ("GeoZ", "Height", "Pressure"),
+        "Y": ("GeoY", "Lat"),
+        "X": ("GeoX", "Lon"),
+    },
+    "axis": {"T": ("T",), "Z": ("Z",), "Y": ("Y",), "X": ("X",)},
+    "positive": {"Z": ("up", "down")},
+    "units": {
+        "Y": (
+            "degree_north",
+            "degree_N",
+            "degreeN",
+            "degrees_north",
+            "degrees_N",
+            "degreesN",
+        ),
+        "X": (
+            "degree_east",
+            "degree_E",
+            "degreeE",
+            "degrees_east",
+            "degrees_E",
+            "degreesE",
+        ),
+    },
+    # "regular_expression": {
+    #     "time": r"time[0-9]*",
+    #     "vertical": (
+    #         r"(lv_|bottom_top|sigma|h(ei)?ght|altitude|depth|isobaric|pres|"
+    #         r"isotherm)[a-z_]*[0-9]*"
+    #     ),
+    #     "y": r"y",
+    #     "latitude": r"x?lat[a-z0-9]*",
+    #     "x": r"x",
+    #     "longitude": r"x?lon[a-z0-9]*",
+    # },
+}
+
+
+def _get_axis_coord(var: xr.DataArray, key, error: bool = True, default: Any = None):
+    """
+    Translate from axis or coord name to variable name
 
+    Parameters
+    ----------
+    var : `xarray.DataArray`
+        DataArray belonging to the coordinate to be checked
+    key : str, ["X", "Y", "Z", "T", "longitude", "latitude", "vertical", "time"]
+        key to check for.
+    error : bool
+        raise errors when key is not found or interpretable. Use False and provide default
+        to replicate dict.get(k, None).
+    default: Any
+        default value to return when error is False.
+
+    Returns
+    -------
+    str, Variable name in parent xarray object that matches axis or coordinate `key`
+
+    Notes
+    -----
+    This functions checks for the following attributes in order
+       - `standard_name` (CF option)
+       - `_CoordinateAxisType` (from THREDDS)
+       - `axis` (CF option)
+       - `positive` (CF standard for non-pressure vertical coordinate)
+
+    References
+    ----------
+    MetPy's parse_cf
+    """
+
+    axis = None
+    if key in _COORD_NAMES:
+        coord = key
+        axis = _COORD_AXIS_MAPPING[key]
+    elif key in _AXIS_NAMES:
+        coord = ""
+        axis = key
+    else:
+        if error:
+            raise KeyError(f"Did not understand {key}")
+        else:
+            return default
+
+    if axis is None:
+        raise AssertionError("Should be unreachable")
+
+    for coord in var.coords:
+        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
+
+    if error:
+        raise KeyError(f"axis name {key!r} not found!")
+    else:
+        return default
+
+
+def _get_measure(da: xr.DataArray, key: str):
+    """
+    TODO: actually interpret da.attrs to get this.
+    """
+    if key not in _CELL_MEASURES:
+        raise ValueError(
+            f"Cell measure must be one of {_CELL_MEASURES!r}. Received {key!r} instead."
+        )
 
-def _get_axis_name_mapping(da: xr.DataArray):
-    return {"X": "lon", "Y": "lat", "T": "time"}
+    return {"area": "cell_area", "volume": "cell_volume"}
 
 
 def _getattr(
@@ -98,12 +237,9 @@ def __getattr__(self, attr):
         )
 
 
-@xr.register_dataarray_accessor("cf")
-@xr.register_dataset_accessor("cf")
 class CFAccessor:
     def __init__(self, da):
         self._obj = da
-        self._coords = _get_axis_name_mapping(da)
 
     def _process_signature(self, func, args, kwargs, keys):
         sig = inspect.signature(func, follow_wrapped=False)
@@ -145,12 +281,17 @@ def _rewrite_values_with_axis_names(self, kwargs, keys, var_kws):
 
                 if isinstance(value, dict):
                     # this for things like isel where **kwargs captures things like T=5
-                    updates[key] = {self._coords.get(k, k): v for k, v in value.items()}
+                    updates[key] = {
+                        _get_axis_coord(self._obj, k, False, k): v
+                        for k, v in value.items()
+                    }
                 elif value is Ellipsis:
                     pass
                 else:
                     # things like sum which have dim
-                    updates[key] = [self._coords.get(v, v) for v in value]
+                    updates[key] = [
+                        _get_axis_coord(self._obj, v, False, v) for v in value
+                    ]
                     if len(updates[key]) == 1:
                         updates[key] = updates[key][0]
 
@@ -163,7 +304,7 @@ def _rewrite_values_with_axis_names(self, kwargs, keys, var_kws):
         for vkw in var_kws:
             if vkw in kwargs:
                 maybe_update = {
-                    k: self._coords.get(v, v)
+                    k: _get_axis_coord(self._obj, v, False, v)
                     for k, v in kwargs[vkw].items()
                     if k in keys
                 }
@@ -177,3 +318,30 @@ def __getattr__(self, attr):
     @property
     def plot(self):
         return _CFWrappedPlotMethods(self._obj, self)
+
+
+@xr.register_dataset_accessor("cf")
+class CFDatasetAccessor(CFAccessor):
+    def __getitem__(self, key):
+        if key in _AXIS_NAMES + _COORD_NAMES:
+            return self._obj[_get_axis_coord(self._obj, key)]
+        elif key in _CELL_MEASURES:
+            raise NotImplementedError("measures not implemented 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)]
+        elif key in _CELL_MEASURES:
+            raise NotImplementedError("measures not implemented yet.")
+            # return self._obj[_get_measure(self._obj)[key]]
+        else:
+            raise KeyError(f"DataArray.cf does not understand the key {key}")

cf_xarray/tests/test_accessor.py

@@ -100,3 +100,30 @@ def test_dataarray_plot(obj):
 @pytest.mark.parametrize("obj", datasets)
 def test_dataset_plot(obj):
     pass
+
+
+@pytest.mark.parametrize("obj", objects)
+@pytest.mark.parametrize(
+    "key, expected_key",
+    (
+        ("X", "lon"),
+        ("Y", "lat"),
+        ("T", "time"),
+        ("longitude", "lon"),
+        ("latitude", "lat"),
+        ("time", "time"),
+    ),
+)
+def test_getitem(obj, key, expected_key):
+    actual = obj.cf[key]
+    expected = obj[expected_key]
+    assert_identical(actual, expected)
+
+
+@pytest.mark.parametrize("obj", objects)
+def test_getitem_errors(obj,):
+    with pytest.raises(KeyError):
+        obj.cf["XX"]
+    obj.lon.attrs = {}
+    with pytest.raises(KeyError):
+        obj.cf["X"]