Support GroupBy iteration and arithmetic (#158)

Fixes #157

Author: dcherian

cf_xarray/accessor.py

@@ -19,6 +19,7 @@
 
 import xarray as xr
 from xarray import DataArray, Dataset
+from xarray.core.arithmetic import SupportsArithmetic
 
 from .helpers import bounds_to_vertices
 from .utils import _is_datetime_like, invert_mappings, parse_cell_methods_attr
@@ -693,7 +694,7 @@ def _get_possible(accessor, criteria):
         return _get_possible(obj.cf, y_criteria)
 
 
-class _CFWrappedClass:
+class _CFWrappedClass(SupportsArithmetic):
     """
     This class is used to wrap any class in _WRAPPED_CLASSES.
     """
@@ -721,6 +722,9 @@ def __getattr__(self, attr):
             key_mappers=_DEFAULT_KEY_MAPPERS,
         )
 
+    def __iter__(self):
+        return iter(self.wrapped)
+
 
 class _CFWrappedPlotMethods:
     """

cf_xarray/tests/test_accessor.py

@@ -786,3 +786,18 @@ def makeds(*dims):
     xtds = makeds("X", "T")
     assert _possible_x_y_plot(xtds, "y") is None
     assert _possible_x_y_plot(xtds, "x") == "X"
+
+
+def test_groupby_special_ops():
+    cfgrouped = airds.cf.groupby_bins("latitude", np.arange(20, 50, 10))
+    grouped = airds.groupby_bins("lat", np.arange(20, 50, 10))
+
+    # __iter__
+    for (label, group), (cflabel, cfgroup) in zip(grouped, cfgrouped):
+        assert label == cflabel
+        assert_identical(group, cfgroup)
+
+    # arithmetic
+    expected = grouped - grouped.mean()
+    actual = grouped - cfgrouped.mean()
+    assert_identical(expected, actual)

doc/whats-new.rst

@@ -15,6 +15,7 @@ v0.4.1 (unreleased)
   by :py:meth:`Dataset.cf.__getitem__`. This allows extraction of DataArrays when there are clashes
   between DataArray names and "special" CF names like ``T``.
   (:issue:`129`, :pr:`130`). By `Deepak Cherian`_
+- Fix iteration and arithemtic with ``GroupBy`` objects. By `Deepak Cherian`_.
 
 v0.4.0 (Jan 22, 2021)
 =====================