Support dictionary-style CF indexing (#661)

Author: Kkkakania

cf_xarray/accessor.py

@@ -1406,6 +1406,31 @@ def __iter__(self):
         return iter(self.wrapped)
 
 
+class _CFWrappedLocIndexer:
+    """Wrap xarray's label-based indexer with CF name translation."""
+
+    def __init__(self, accessor: CFAccessor):
+        self.accessor = accessor
+
+    def _translate_indexers(self, key):
+        if not isinstance(key, Mapping):
+            return key
+
+        _, arguments = self.accessor._process_signature(
+            self.accessor._obj.sel,
+            (key,),
+            {},
+            dict(_DEFAULT_KEY_MAPPERS),
+        )
+        return arguments["indexers"]
+
+    def __getitem__(self, key):
+        return self.accessor._obj.loc[self._translate_indexers(key)]
+
+    def __setitem__(self, key, value):
+        self.accessor._obj.loc[self._translate_indexers(key)] = value
+
+
 class _CFWrappedPlotMethods:
     """
     This class wraps DataArray.plot
@@ -1942,6 +1967,11 @@ def __contains__(self, item: str) -> bool:
         """
         return item in self.keys()
 
+    @property
+    def loc(self):
+        """Label-based indexer that understands CF names."""
+        return _CFWrappedLocIndexer(self)
+
     @property
     def plot(self):
         """
@@ -2809,13 +2839,15 @@ def grid_mappings(self) -> tuple[GridMapping, ...]:
 
 @xr.register_dataset_accessor("cf")
 class CFDatasetAccessor(CFAccessor):
-    def __getitem__(self, key: Hashable | Iterable[Hashable]) -> DataArray | Dataset:
+    def __getitem__(
+        self, key: Hashable | Iterable[Hashable] | Mapping[Hashable, Any]
+    ) -> DataArray | Dataset:
         """
         Index into a Dataset making use of CF attributes.
 
         Parameters
         ----------
-        key : str, Iterable[str], optional
+        key : str, Iterable[str], or Mapping, optional
             One of
               - axes names: "X", "Y", "Z", "T"
               - coordinate names: "longitude", "latitude", "vertical", "time"
@@ -2841,6 +2873,9 @@ def __getitem__(self, key: Hashable | Iterable[Hashable]) -> DataArray | Dataset
 
         Add additional keys by specifying "custom criteria". See :ref:`custom_criteria` for more.
         """
+        if isinstance(key, Mapping):
+            return self.isel(key)
+
         return _getitem(self, key)
 
     @property
@@ -3474,13 +3509,15 @@ def grid_mapping_name(self) -> str:
         # Return the single grid mapping name
         return next(iter(grid_mapping_names.keys()))
 
-    def __getitem__(self, key: Hashable | Iterable[Hashable]) -> DataArray:
+    def __getitem__(
+        self, key: Hashable | Iterable[Hashable] | Mapping[Hashable, Any]
+    ) -> DataArray:
         """
         Index into a DataArray making use of CF attributes.
 
         Parameters
         ----------
-        key : str, Iterable[str], optional
+        key : str, Iterable[str], or Mapping, optional
             One of
               - axes names: "X", "Y", "Z", "T"
               - coordinate names: "longitude", "latitude", "vertical", "time"
@@ -3509,6 +3546,9 @@ def __getitem__(self, key: Hashable | Iterable[Hashable]) -> DataArray:
         Add additional keys by specifying "custom criteria". See :ref:`custom_criteria` for more.
         """
 
+        if isinstance(key, Mapping):
+            return self.isel(key)
+
         if not isinstance(key, Hashable):
             raise KeyError(
                 f"Cannot use an Iterable of keys with DataArrays. Expected a single string. Received {key!r} instead."

cf_xarray/tests/test_accessor.py

@@ -517,6 +517,42 @@ def test_kwargs_methods(obj):
     assert_identical(expected, actual)
 
 
+@pytest.mark.parametrize("obj", objects)
+def test_dictionary_indexing(obj):
+    with raise_if_dask_computes():
+        expected = obj.isel(time=0)
+        actual = obj.cf[{"T": 0}]
+    assert_identical(expected, actual)
+
+
+@pytest.mark.parametrize("obj", objects)
+def test_loc_indexing(obj):
+    label = obj.time.values[0]
+    with raise_if_dask_computes():
+        expected = obj.loc[{"time": label}]
+        actual = obj.cf.loc[{"T": label}]
+    assert_identical(expected, actual)
+
+
+@pytest.mark.parametrize("obj", objects)
+def test_loc_assignment(obj):
+    expected = obj.copy()
+    actual = obj.copy()
+    label = obj.time.values[0]
+
+    expected.loc[{"time": label}] = -1
+    actual.cf.loc[{"T": label}] = -1
+
+    assert_identical(expected, actual)
+
+
+def test_dictionary_indexing_expands_multiple_dimensions():
+    expected = multiple.isel(x1=5, x2=5)
+    actual = multiple.cf[{"X": 5}]
+
+    assert_identical(expected, actual)
+
+
 def test_pos_args_methods() -> None:
     expected = airds.transpose("lon", "time", "lat")
     actual = airds.cf.transpose("longitude", "T", "latitude")