Support dictionary-style indexing with .loc

[lukelbd]

Given the following sample data:

import numpy as np
import xarray as xr
import cf_xarray
lev = np.arange(1, 11)
lev = xr.DataArray(lev, dims='lev', attrs={'units': 'km', 'axis': 'Z'})
data = np.arange(lev.size)
data = xr.DataArray(data, dims='lev', coords={'lev': lev})
print(data.cf.axes)  # prints {'Z': ['lev']}
print(data[{'lev': 0}])  # prints data for level value "1"
print(data.loc[{'lev': 1}])  # prints data for level value "1"

The cf accessor cannot handle dictionary-style position-based indexing:

>>> data.cf[{'Z': 0}]
File ~/miniconda3/lib/python3.10/site-packages/cf_xarray/accessor.py:2402, in CFDataArrayAccessor.__getitem__(self, key)
   2370 """
   2371 Index into a DataArray making use of CF attributes.
   2372
   (...)
   2398 Add additional keys by specifying "custom criteria". See :ref:`custom_criteria` for more.
   2399 """
   2401 if not isinstance(key, str):
-> 2402     raise KeyError(
   2403         f"Cannot use a list of keys with DataArrays. Expected a single string. Received {key!r} instead."
   2404     )
   2406 return _getitem(self, key)

KeyError: "Cannot use a list of keys with DataArrays. Expected a single string. Received {'Z': 0} instead."

It also cannot handle dictionary-style coordinate-based inexing:

>>> data.cf.loc[{'Z': 0}]
File ~/miniconda3/lib/python3.10/site-packages/cf_xarray/accessor.py:576, in _getattr(obj, attr, accessor, key_mappers, wrap_classes, extra_decorator)
    573     func: Callable = attribute
    575 else:
--> 576     raise AttributeError(
    577         f"cf_xarray does not know how to wrap attribute '{type(obj).__name__}.{attr}'. "
    578         "Please file an issue if you have a solution."
    579     )
    581 @functools.wraps(func)
    582 def wrapper(*args, **kwargs):
    583     posargs, arguments = accessor._process_signature(
    584         func, args, kwargs, key_mappers
    585     )

AttributeError: cf_xarray does not know how to wrap attribute 'DataArray.loc'. Please file an issue if you have a solution.

I'd be interested in addressing this but don't currently have time. Thought I'd just get the issue on this project's radar for now (if it wasn't already / I didn't miss an already-open thread).

[dcherian]

👍 should be straightforward. I don't use this syntax, that's why it hasn't been implemented yet :)