diff --git a/doc/whats-new.rst b/doc/whats-new.rst index f9c952f6752..7103d7b8ab3 100644 --- a/doc/whats-new.rst +++ b/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 `_. + 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 diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index bae3057aabe..3399b27b3e7 100644 --- a/xarray/core/groupby.py +++ b/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 diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index ee17cc39064..463a6a32185 100644 --- a/xarray/tests/test_groupby.py +++ b/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