Rework getitem (#65)

1. If passes a scalar key, there must be only one variable that matches
the key
2. To get back all variables, pass a list ds.cf[["longitude"]] for
example.

I think this is more consistent with xarray semantics.

Author: dcherian

cf_xarray/accessor.py

@@ -174,11 +174,11 @@ def _get_axis_coord_single(var: Union[DataArray, Dataset], key: str,) -> List[st
     """ Helper method for when we really want only one result per key. """
     results = _get_axis_coord(var, key)
     if len(results) > 1:
-        raise ValueError(
+        raise KeyError(
             f"Multiple results for {key!r} found: {results!r}. Is this valid CF? Please open an issue."
         )
     elif len(results) == 0:
-        raise ValueError(f"No results found for {key!r}.")
+        raise KeyError(f"No results found for {key!r}.")
     return results
 
 
@@ -895,14 +895,17 @@ def __getitem__(self, key: Union[str, List[str]]):
             )
 
         if scalar_key:
+            axis_coord_mapper = _get_axis_coord_single
             key = (key,)  # type: ignore
+        else:
+            axis_coord_mapper = _get_axis_coord
 
         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)
+                names = axis_coord_mapper(self._obj, k)
                 successful[k] = bool(names)
                 coords.extend(names)
             elif k in _CELL_MEASURES:
@@ -919,6 +922,7 @@ def __getitem__(self, key: Union[str, List[str]]):
         # 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 varnames:
@@ -929,8 +933,10 @@ def __getitem__(self, key: Union[str, List[str]]):
             else:
                 ds = self._obj
 
-            if scalar_key and len(varnames) == 1:
-                da: DataArray = ds[varnames[0]].reset_coords(drop=True)  # type: ignore
+            if scalar_key and len(allnames) == 1:
+                da: DataArray = ds.reset_coords()[allnames[0]]  # type: ignore
+                if allnames[0] in coords:
+                    coords.remove(allnames[0])
                 failed = []
                 for k1 in coords:
                     if k1 not in ds.variables:

cf_xarray/tests/test_accessor.py

@@ -105,7 +105,7 @@ def test_getitem_ancillary_variables():
 def test_rename_like():
     original = popds.copy(deep=True)
 
-    with pytest.raises(ValueError):
+    with pytest.raises(KeyError):
         popds.cf.rename_like(airds)
 
     renamed = popds.cf["TEMP"].cf.rename_like(airds)
@@ -310,8 +310,6 @@ def test_getitem(obj, key, expected_key):
     assert key in obj.cf
 
     actual = obj.cf[key]
-    if isinstance(obj, xr.Dataset):
-        expected_key = [expected_key]
     expected = obj[expected_key]
     assert_identical(actual, expected)
 
@@ -329,10 +327,11 @@ def test_getitem_uses_coordinates():
     # POP-like dataset
     ds = popds
     assert_identical(
-        ds.cf["X"], ds.reset_coords()[["ULONG", "TLONG"]].set_coords(["ULONG", "TLONG"])
+        ds.cf[["X"]],
+        ds.reset_coords()[["ULONG", "TLONG"]].set_coords(["ULONG", "TLONG"]),
     )
     assert_identical(
-        ds.cf["Y"], ds.reset_coords()[["ULAT", "TLAT"]].set_coords(["ULAT", "TLAT"])
+        ds.cf[["Y"]], ds.reset_coords()[["ULAT", "TLAT"]].set_coords(["ULAT", "TLAT"])
     )
     assert_identical(ds.UVEL.cf["X"], ds["ULONG"].reset_coords(drop=True))
     assert_identical(ds.TEMP.cf["X"], ds["TLONG"].reset_coords(drop=True))

doc/contributing.rst

@@ -55,7 +55,6 @@ Primarily for developer reference. Some of these could become public API if nece
     ~accessor._getattr
     ~accessor._get_axis_coord
     ~accessor._get_axis_coord_single
-    ~accessor._get_list_standard_names
     ~accessor._get_measure
     ~accessor._get_measure_variable
     accessor._CFWrappedPlotMethods._plot_decorator