update

Author: dcherian

cf_xarray/groupers.py

@@ -1,27 +1,33 @@
+from dataclasses import dataclass
+
 import numpy as np
 import pandas as pd
-from xarray.groupers import EncodedGroups, Grouper
+from xarray.groupers import EncodedGroups, UniqueGrouper
 
 
-class FlagGrouper(Grouper):
+@dataclass
+class FlagGrouper(UniqueGrouper):
     def factorize(self, group) -> EncodedGroups:
-        assert "flag_values" in group.attrs
-        assert "flag_meanings" in group.attrs
+        if "flag_values" not in group.attrs or "flag_meanings" not in group.attrs:
+            raise ValueError(
+                "FlagGrouper can only be used with flag variables that have"
+                "`flag_values` and `flag_meanings` specified in attrs."
+            )
 
         values = np.array(group.attrs["flag_values"])
         full_index = pd.Index(group.attrs["flag_meanings"].split(" "))
 
-        if group.dtype.kind in "iu" and (np.diff(values) == 1).all():
-            # optimize
-            codes = group.data - values[0].astype(int)
-        else:
-            codes, _ = pd.factorize(group.data.ravel())
+        self.labels = values
+        ret = super().factorize(group)
 
-        codes_da = group.copy(data=codes.reshape(group.shape))
+        codes_da = ret.codes
         codes_da.attrs.pop("flag_values")
         codes_da.attrs.pop("flag_meanings")
 
-        return EncodedGroups(codes=codes_da, full_index=full_index)
+        ret.codes = codes_da
+        ret.full_index = full_index
+
+        return ret
 
     def reset(self):
         pass

cf_xarray/tests/test_groupers.py

@@ -1,15 +1,45 @@
 import numpy as np
+import pytest
+import xarray as xr
 from xarray.testing import assert_identical
 
 from cf_xarray.datasets import flag_excl
 from cf_xarray.groupers import FlagGrouper
 
 
 def test_flag_grouper():
-    ds = flag_excl.to_dataset().set_coords("flag_var")
+    ds = flag_excl.to_dataset().set_coords("flag_var").copy(deep=True)
     ds["foo"] = ("time", np.arange(8))
     actual = ds.groupby(flag_var=FlagGrouper()).mean()
     expected = ds.groupby("flag_var").mean()
     expected["flag_var"] = ["flag_1", "flag_2", "flag_3"]
     expected["flag_var"].attrs["standard_name"] = "flag_mutual_exclusive"
     assert_identical(actual, expected)
+
+    del ds.flag_var.attrs["flag_values"]
+    with pytest.raises(ValueError):
+        ds.groupby(flag_var=FlagGrouper())
+
+    ds.flag_var.attrs["flag_values"] = [0, 1, 2]
+    del ds.flag_var.attrs["flag_meanings"]
+    with pytest.raises(ValueError):
+        ds.groupby(flag_var=FlagGrouper())
+
+
+@pytest.mark.parametrize(
+    "values",
+    [
+        [1, 2],
+        [1, 2, 3],  # value out of range of flag_values
+    ],
+)
+def test_flag_grouper_optimized(values):
+    ds = xr.Dataset(
+        {"foo": ("x", values, {"flag_values": [0, 1, 2], "flag_meanings": "a b c"})}
+    )
+    ret = FlagGrouper().factorize(ds.foo)
+    expected = ds.foo
+    expected.data[ds.foo.data > 2] = -1
+    del ds.foo.attrs["flag_meanings"]
+    del ds.foo.attrs["flag_values"]
+    assert_identical(ret.codes, ds.foo)