Skip to content

Commit 86fb71d

Browse files
authored
groupby repr (#3344)
* groupby repr * add test. * test datetime and nondim coord * rename test_da → repr_da * Add whats-new * Update doc/whats-new.rst
1 parent dd2b803 commit 86fb71d

File tree

3 files changed

+71
-0
lines changed

3 files changed

+71
-0
lines changed

doc/whats-new.rst

+13
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,19 @@ What's New
1818
v0.13.1 (unreleased)
1919
--------------------
2020

21+
New functions/methods
22+
~~~~~~~~~~~~~~~~~~~~~
23+
24+
Enhancements
25+
~~~~~~~~~~~~
26+
27+
- Add a repr for :py:class:`~xarray.core.GroupBy` objects. By `Deepak Cherian <https://github.com/dcherian>`_.
28+
Example::
29+
30+
>>> da.groupby("time.season")
31+
DataArrayGroupBy, grouped over 'season'
32+
4 groups with labels 'DJF', 'JJA', 'MAM', 'SON'
33+
2134
Bug fixes
2235
~~~~~~~~~
2336
- Reintroduce support for :mod:`weakref` (broken in v0.13.0). Support has been

xarray/core/groupby.py

+18
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from .arithmetic import SupportsArithmetic
1010
from .common import ImplementsArrayReduce, ImplementsDatasetReduce
1111
from .concat import concat
12+
from .formatting import format_array_flat
1213
from .options import _get_keep_attrs
1314
from .pycompat import integer_types
1415
from .utils import (
@@ -158,6 +159,15 @@ def ndim(self):
158159
def values(self):
159160
return range(self.size)
160161

162+
@property
163+
def shape(self):
164+
return (self.size,)
165+
166+
def __getitem__(self, key):
167+
if isinstance(key, tuple):
168+
key = key[0]
169+
return self.values[key]
170+
161171

162172
def _ensure_1d(group, obj):
163173
if group.ndim != 1:
@@ -383,6 +393,14 @@ def __len__(self):
383393
def __iter__(self):
384394
return zip(self._unique_coord.values, self._iter_grouped())
385395

396+
def __repr__(self):
397+
return "%s, grouped over %r \n%r groups with labels %s" % (
398+
self.__class__.__name__,
399+
self._unique_coord.name,
400+
self._unique_coord.size,
401+
", ".join(format_array_flat(self._unique_coord, 30).split()),
402+
)
403+
386404
def _get_index_and_items(self, index, grouper):
387405
from .resample_cftime import CFTimeGrouper
388406

xarray/tests/test_groupby.py

+40
Original file line numberDiff line numberDiff line change
@@ -202,4 +202,44 @@ def test_da_groupby_assign_coords():
202202
assert_identical(expected, actual2)
203203

204204

205+
repr_da = xr.DataArray(
206+
np.random.randn(10, 20, 6, 24),
207+
dims=["x", "y", "z", "t"],
208+
coords={
209+
"z": ["a", "b", "c", "a", "b", "c"],
210+
"x": [1, 1, 1, 2, 2, 3, 4, 5, 3, 4],
211+
"t": pd.date_range("2001-01-01", freq="M", periods=24),
212+
"month": ("t", list(range(1, 13)) * 2),
213+
},
214+
)
215+
216+
217+
@pytest.mark.parametrize("dim", ["x", "y", "z", "month"])
218+
@pytest.mark.parametrize("obj", [repr_da, repr_da.to_dataset(name="a")])
219+
def test_groupby_repr(obj, dim):
220+
actual = repr(obj.groupby(dim))
221+
expected = "%sGroupBy" % obj.__class__.__name__
222+
expected += ", grouped over %r " % dim
223+
expected += "\n%r groups with labels " % (len(np.unique(obj[dim])))
224+
if dim == "x":
225+
expected += "1, 2, 3, 4, 5"
226+
elif dim == "y":
227+
expected += "0, 1, 2, 3, 4, 5, ..., 15, 16, 17, 18, 19"
228+
elif dim == "z":
229+
expected += "'a', 'b', 'c'"
230+
elif dim == "month":
231+
expected += "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12"
232+
assert actual == expected
233+
234+
235+
@pytest.mark.parametrize("obj", [repr_da, repr_da.to_dataset(name="a")])
236+
def test_groupby_repr_datetime(obj):
237+
actual = repr(obj.groupby("t.month"))
238+
expected = "%sGroupBy" % obj.__class__.__name__
239+
expected += ", grouped over 'month' "
240+
expected += "\n%r groups with labels " % (len(np.unique(obj.t.dt.month)))
241+
expected += "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12"
242+
assert actual == expected
243+
244+
205245
# TODO: move other groupby tests from test_dataset and test_dataarray over here

0 commit comments

Comments
 (0)