diff --git a/cf_xarray/accessor.py b/cf_xarray/accessor.py index ed0c216d..7b0a09d9 100644 --- a/cf_xarray/accessor.py +++ b/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): """ @@ -2803,13 +2833,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" @@ -2835,6 +2867,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 @@ -3468,13 +3503,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" @@ -3503,6 +3540,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." diff --git a/cf_xarray/tests/test_accessor.py b/cf_xarray/tests/test_accessor.py index 6657be46..1197a9d9 100644 --- a/cf_xarray/tests/test_accessor.py +++ b/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")