Rework attribute rewriting, .cf.data_vars, .cf.coords (#130)

* Fix .cf.data_vars
* Return "CF DataArrays" in .cf.data_vars,.cf.coords
* Fix DataArray.cf["standard_name"]

Closes #129
Closes #126

Author: dcherian

cf_xarray/accessor.py

@@ -2,7 +2,7 @@
 import inspect
 import itertools
 import warnings
-from collections import ChainMap
+from collections import ChainMap, defaultdict
 from typing import (
     Any,
     Callable,
@@ -159,6 +159,18 @@ def _is_datetime_like(da: DataArray) -> bool:
     return False
 
 
+def invert_mappings(*mappings):
+    """Takes a set of mappings and iterates through, inverting to make a
+    new mapping of value: set(keys). Keys are deduplicated to avoid clashes between
+    standard_name and coordinate names."""
+    merged = defaultdict(set)
+    for mapping in mappings:
+        for k, v in mapping.items():
+            for name in v:
+                merged[name] |= set([k])
+    return merged
+
+
 # Type for Mapper functions
 Mapper = Callable[[Union[DataArray, Dataset], str], List[str]]
 
@@ -503,23 +515,29 @@ def _getattr(
     if isinstance(attribute, Mapping):
         if not attribute:
             return dict(attribute)
-        # attributes like chunks / sizes
+
         newmap = dict()
-        unused_keys = set(attribute.keys())
-        for key in _AXIS_NAMES + _COORD_NAMES:
-            value = set(apply_mapper(_get_axis_coord, obj, key, error=False))
-            unused_keys -= value
-            if value:
-                good_values = value & set(obj.dims)
-                if not good_values:
-                    continue
-                if len(good_values) > 1:
+        inverted = invert_mappings(
+            accessor.axes,
+            accessor.coordinates,
+            accessor.cell_measures,
+            accessor.standard_names,
+        )
+        unused_keys = set(attribute.keys()) - set(inverted)
+        for key, value in attribute.items():
+            for name in inverted[key]:
+                if name in newmap:
                     raise AttributeError(
-                        f"cf_xarray can't wrap attribute {attr!r} because there are multiple values for {key!r} viz. {good_values!r}. "
-                        f"There is no unique mapping from {key!r} to a value in {attr!r}."
+                        f"cf_xarray can't wrap attribute {attr!r} because there are multiple values for {name!r}. "
+                        f"There is no unique mapping from {name!r} to a value in {attr!r}."
                     )
-                newmap.update({key: attribute[good_values.pop()]})
+            newmap.update(dict.fromkeys(inverted[key], value))
         newmap.update({key: attribute[key] for key in unused_keys})
+
+        skip = {"data_vars": ["coords"], "coords": None}
+        if attr in ["coords", "data_vars"]:
+            for key in newmap:
+                newmap[key] = _getitem(accessor, key, skip=skip[attr])
         return newmap
 
     elif isinstance(attribute, Callable):  # type: ignore
@@ -548,6 +566,123 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def _getitem(
+    accessor: "CFAccessor", key: Union[str, List[str]], skip: List[str] = None
+) -> Union[DataArray, Dataset]:
+    """
+    Index into obj using key. Attaches CF associated variables.
+
+    Parameters
+    ----------
+    accessor: CFAccessor
+    key: str, List[str]
+    skip: str, optional
+        One of ["coords", "measures"], avoid clashes with special coord names
+    """
+
+    obj = accessor._obj
+    kind = str(type(obj).__name__)
+    scalar_key = isinstance(key, str)
+
+    if isinstance(obj, DataArray) and not scalar_key:
+        raise KeyError(
+            f"Cannot use a list of keys with DataArrays. Expected a single string. Received {key!r} instead."
+        )
+
+    if scalar_key:
+        key = (key,)  # type: ignore
+
+    if skip is None:
+        skip = []
+
+    def check_results(names, k):
+        if scalar_key and len(names) > 1:
+            raise ValueError(
+                f"Receive multiple variables for key {k!r}: {names}. "
+                f"Expected only one. Please pass a list [{k!r}] "
+                f"instead to get all variables matching {k!r}."
+            )
+
+    varnames: List[Hashable] = []
+    coords: List[Hashable] = []
+    successful = dict.fromkeys(key, False)
+    for k in key:
+        if "coords" not in skip and k in _AXIS_NAMES + _COORD_NAMES:
+            names = _get_axis_coord(obj, k)
+            check_results(names, k)
+            successful[k] = bool(names)
+            coords.extend(names)
+        elif "measures" not in skip and k in accessor._get_all_cell_measures():
+            measure = _get_measure(obj, k)
+            check_results(measure, k)
+            successful[k] = bool(measure)
+            if measure:
+                varnames.extend(measure)
+        else:
+            stdnames = set(_get_with_standard_name(obj, k))
+            check_results(stdnames, k)
+            successful[k] = bool(stdnames)
+            objcoords = set(obj.coords)
+            varnames.extend(stdnames - objcoords)
+            coords.extend(stdnames & objcoords)
+
+    # these are not special names but could be variable names in underlying object
+    # we allow this so that we can return variables with appropriate CF auxiliary variables
+    varnames.extend([k for k, v in successful.items() if not v])
+    allnames = varnames + coords
+
+    try:
+        for name in allnames:
+            extravars = accessor.get_associated_variable_names(name)
+            # we cannot return bounds variables with scalar keys
+            if scalar_key:
+                extravars.pop("bounds")
+            coords.extend(itertools.chain(*extravars.values()))
+
+        if isinstance(obj, DataArray):
+            ds = obj._to_temp_dataset()
+        else:
+            ds = obj
+
+        if scalar_key:
+            if len(allnames) == 1:
+                da: DataArray = ds.reset_coords()[allnames[0]]  # type: ignore
+                if allnames[0] in coords:
+                    coords.remove(allnames[0])
+                for k1 in coords:
+                    da.coords[k1] = ds.variables[k1]
+                return da
+            else:
+                raise ValueError(
+                    f"Received scalar key {key[0]!r} but multiple results: {allnames!r}. "
+                    f"Please pass a list instead (['{key[0]}']) to get back a Dataset "
+                    f"with {allnames!r}."
+                )
+
+        ds = ds.reset_coords()[varnames + coords]
+        if isinstance(obj, DataArray):
+            if scalar_key and len(ds.variables) == 1:
+                # single dimension coordinates
+                assert coords
+                assert not varnames
+
+                return ds[coords[0]]
+
+            elif scalar_key and len(ds.variables) > 1:
+                raise NotImplementedError(
+                    "Not sure what to return when given scalar key for DataArray and it has multiple values. "
+                    "Please open an issue."
+                )
+
+        return ds.set_coords(coords)
+
+    except KeyError:
+        raise KeyError(
+            f"{kind}.cf does not understand the key {k!r}. "
+            f"Use {kind}.cf.describe() to see a list of key names that can be interpreted."
+        )
+
+
 class _CFWrappedClass:
     """
     This class is used to wrap any class in _WRAPPED_CLASSES.
@@ -1061,104 +1196,7 @@ def get_associated_variable_names(self, name: Hashable) -> Dict[str, List[str]]:
         return coords
 
     def __getitem__(self, key: Union[str, List[str]]):
-
-        kind = str(type(self._obj).__name__)
-        scalar_key = isinstance(key, str)
-
-        if isinstance(self._obj, DataArray) and not scalar_key:
-            raise KeyError(
-                f"Cannot use a list of keys with DataArrays. Expected a single string. Received {key!r} instead."
-            )
-
-        if scalar_key:
-            key = (key,)  # type: ignore
-
-        def check_results(names, k):
-            if scalar_key and len(names) > 1:
-                raise ValueError(
-                    f"Receive multiple variables for key {k!r}: {names}. "
-                    f"Expected only one. Please pass a list [{k!r}] "
-                    f"instead to get all variables matching {k!r}."
-                )
-
-        varnames: List[Hashable] = []
-        coords: List[Hashable] = []
-        successful = dict.fromkeys(key, False)
-        for k in key:
-            if k in _AXIS_NAMES + _COORD_NAMES:
-                names = _get_axis_coord(self._obj, k)
-                check_results(names, k)
-                successful[k] = bool(names)
-                coords.extend(names)
-            elif k in self._get_all_cell_measures():
-                measure = _get_measure(self._obj, k)
-                check_results(measure, k)
-                successful[k] = bool(measure)
-                if measure:
-                    varnames.extend(measure)
-            elif not isinstance(self._obj, DataArray):
-                stdnames = set(_get_with_standard_name(self._obj, k))
-                check_results(stdnames, k)
-                successful[k] = bool(stdnames)
-                objcoords = set(self._obj.coords)
-                varnames.extend(stdnames - objcoords)
-                coords.extend(stdnames & objcoords)
-
-        # these are not special names but could be variable names in underlying object
-        # we allow this so that we can return variables with appropriate CF auxiliary variables
-        varnames.extend([k for k, v in successful.items() if not v])
-        allnames = varnames + coords
-
-        try:
-            for name in allnames:
-                extravars = self.get_associated_variable_names(name)
-                # we cannot return bounds variables with scalar keys
-                if scalar_key:
-                    extravars.pop("bounds")
-                coords.extend(itertools.chain(*extravars.values()))
-
-            if isinstance(self._obj, DataArray):
-                ds = self._obj._to_temp_dataset()
-            else:
-                ds = self._obj
-
-            if scalar_key:
-                if len(allnames) == 1:
-                    da: DataArray = ds.reset_coords()[allnames[0]]  # type: ignore
-                    if allnames[0] in coords:
-                        coords.remove(allnames[0])
-                    for k1 in coords:
-                        da.coords[k1] = ds.variables[k1]
-                    return da
-                else:
-                    raise ValueError(
-                        f"Received scalar key {key[0]!r} but multiple results: {allnames!r}. "
-                        f"Please pass a list instead (['{key[0]}']) to get back a Dataset "
-                        f"with {allnames!r}."
-                    )
-
-            ds = ds.reset_coords()[varnames + coords]
-            if isinstance(self._obj, DataArray):
-                if scalar_key and len(ds.variables) == 1:
-                    # single dimension coordinates
-                    assert coords
-                    assert not varnames
-
-                    return ds[coords[0]]
-
-                elif scalar_key and len(ds.variables) > 1:
-                    raise NotImplementedError(
-                        "Not sure what to return when given scalar key for DataArray and it has multiple values. "
-                        "Please open an issue."
-                    )
-
-            return ds.set_coords(coords)
-
-        except KeyError:
-            raise KeyError(
-                f"{kind}.cf does not understand the key {k!r}. "
-                f"Use {kind}.cf.describe() to see a list of key names that can be interpreted."
-            )
+        return _getitem(self, key)
 
     def _maybe_to_dataset(self, obj=None) -> Dataset:
         if obj is None:

cf_xarray/tests/test_accessor.py

@@ -19,6 +19,12 @@
 objects = datasets + dataarrays
 
 
+def assert_dicts_identical(dict1, dict2):
+    assert dict1.keys() == dict2.keys()
+    for k in dict1:
+        assert_identical(dict1[k], dict2[k])
+
+
 def test_describe(capsys):
     airds.cf.describe()
     actual = capsys.readouterr().out
@@ -280,7 +286,10 @@ def test_dataarray_getitem():
     with pytest.raises(KeyError):
         air.cf[["longitude"]]
     with pytest.raises(KeyError):
-        air.cf[["longitude", "latitude"]],
+        air.cf[["longitude", "latitude"]]
+
+    air["cell_area"].attrs["standard_name"] = "area_grid_cell"
+    assert_identical(air.cf["area_grid_cell"], air.cell_area.reset_coords(drop=True))
 
 
 @pytest.mark.parametrize("obj", dataarrays)
@@ -512,7 +521,7 @@ def test_guess_coord_axis():
     assert dsnew.y1.attrs == {"axis": "Y"}
 
 
-def test_dicts():
+def test_attributes():
     actual = airds.cf.sizes
     expected = {"X": 50, "Y": 25, "T": 4, "longitude": 50, "latitude": 25, "time": 4}
     assert actual == expected
@@ -543,6 +552,30 @@ def test_dicts():
     expected = {"lon": 50, "Y": 25, "T": 4, "latitude": 25, "time": 4}
     assert actual == expected
 
+    actual = popds.cf.data_vars
+    expected = {
+        "sea_water_x_velocity": popds.cf["UVEL"],
+        "sea_water_potential_temperature": popds.cf["TEMP"],
+    }
+    assert_dicts_identical(actual, expected)
+
+    actual = multiple.cf.data_vars
+    expected = dict(multiple.data_vars)
+    assert_dicts_identical(actual, expected)
+
+    # check that data_vars contains ancillary variables
+    assert_identical(anc.cf.data_vars["specific_humidity"], anc.cf["specific_humidity"])
+
+    # clash between var name and "special" CF name
+    # Regression test for #126
+    data = np.random.rand(4, 3)
+    times = pd.date_range("2000-01-01", periods=4)
+    locs = [30, 60, 90]
+    coords = [("time", times, {"axis": "T"}), ("space", locs)]
+    foo = xr.DataArray(data, coords, dims=["time", "space"])
+    ds1 = xr.Dataset({"T": foo})
+    assert_identical(ds1.cf.data_vars["T"], ds1["T"])
+
 
 def test_missing_variable_in_coordinates():
     airds.air.attrs["coordinates"] = "lat lon time"

doc/whats-new.rst

@@ -6,7 +6,12 @@ What's New
 v0.4.1 (unreleased)
 ===================
 
-- Support for using ``standard_name`` in more functions. By `Deepak Cherian`_
+- Support for using ``standard_name`` in more functions. (:pr:`128`) By `Deepak Cherian`_
+- Allow ``DataArray.cf[]`` with standard names. By `Deepak Cherian`_
+- Rewrite the ``values`` of ``.cf.coords`` and ``.cf.data_vars`` with objects returned
+  by ``.cf.__getitem___``. This allows extraction of DataArrays when there are clashes
+  between DataArray names and "special" CF names like ``T``.
+  (:issue:`129`, :pr:`130`). By `Deepak Cherian`_
 
 v0.4.0 (Jan 22, 2021)
 =====================