groupby repr

doc/whats-new.rst

@@ -18,6 +18,19 @@ What's New
 v0.13.1 (unreleased)
 --------------------
 
+New functions/methods
+~~~~~~~~~~~~~~~~~~~~~
+
+Enhancements
+~~~~~~~~~~~~
+
+- Add a repr for :py:class:`~xarray.core.GroupBy` objects. By `Deepak Cherian <https://github.com/dcherian>`_.
+  Example::
+
+      >>> da.groupby("time.season")
+      DataArrayGroupBy, grouped over 'season'
+      4 groups with labels 'DJF', 'JJA', 'MAM', 'SON
+
 Bug fixes
 ~~~~~~~~~
 - Reintroduce support for :mod:`weakref` (broken in v0.13.0). Support has been

xarray/core/groupby.py

@@ -9,6 +9,7 @@
 from .arithmetic import SupportsArithmetic
 from .common import ImplementsArrayReduce, ImplementsDatasetReduce
 from .concat import concat
+from .formatting import format_array_flat
 from .options import _get_keep_attrs
 from .pycompat import integer_types
 from .utils import (
@@ -158,6 +159,15 @@ def ndim(self):
     def values(self):
         return range(self.size)
 
+    @property
+    def shape(self):
+        return (self.size,)
+
+    def __getitem__(self, key):
+        if isinstance(key, tuple):
+            key = key[0]
+        return self.values[key]
+
 
 def _ensure_1d(group, obj):
     if group.ndim != 1:
@@ -383,6 +393,14 @@ def __len__(self):
     def __iter__(self):
         return zip(self._unique_coord.values, self._iter_grouped())
 
+    def __repr__(self):
+        return "%s, grouped over %r \n%r groups with labels %s" % (
+            self.__class__.__name__,
+            self._unique_coord.name,
+            self._unique_coord.size,
+            ", ".join(format_array_flat(self._unique_coord, 30).split()),
+        )
+
     def _get_index_and_items(self, index, grouper):
         from .resample_cftime import CFTimeGrouper
 

xarray/tests/test_groupby.py

@@ -202,4 +202,44 @@ def test_da_groupby_assign_coords():
     assert_identical(expected, actual2)
 
 
+repr_da = xr.DataArray(
+    np.random.randn(10, 20, 6, 24),
+    dims=["x", "y", "z", "t"],
+    coords={
+        "z": ["a", "b", "c", "a", "b", "c"],
+        "x": [1, 1, 1, 2, 2, 3, 4, 5, 3, 4],
+        "t": pd.date_range("2001-01-01", freq="M", periods=24),
+        "month": ("t", list(range(1, 13)) * 2),
+    },
+)
+
+
+@pytest.mark.parametrize("dim", ["x", "y", "z", "month"])
+@pytest.mark.parametrize("obj", [repr_da, repr_da.to_dataset(name="a")])
+def test_groupby_repr(obj, dim):
+    actual = repr(obj.groupby(dim))
+    expected = "%sGroupBy" % obj.__class__.__name__
+    expected += ", grouped over %r " % dim
+    expected += "\n%r groups with labels " % (len(np.unique(obj[dim])))
+    if dim == "x":
+        expected += "1, 2, 3, 4, 5"
+    elif dim == "y":
+        expected += "0, 1, 2, 3, 4, 5, ..., 15, 16, 17, 18, 19"
+    elif dim == "z":
+        expected += "'a', 'b', 'c'"
+    elif dim == "month":
+        expected += "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12"
+    assert actual == expected
+
+
+@pytest.mark.parametrize("obj", [repr_da, repr_da.to_dataset(name="a")])
+def test_groupby_repr_datetime(obj):
+    actual = repr(obj.groupby("t.month"))
+    expected = "%sGroupBy" % obj.__class__.__name__
+    expected += ", grouped over 'month' "
+    expected += "\n%r groups with labels " % (len(np.unique(obj.t.dt.month)))
+    expected += "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12"
+    assert actual == expected
+
+
 # TODO: move other groupby tests from test_dataset and test_dataarray over here