Skip to content

Commit

Permalink
groupby repr (#3344)
Browse files Browse the repository at this point in the history
* groupby repr

* add test.

* test datetime and nondim coord

* rename test_da → repr_da

* Add whats-new

* Update doc/whats-new.rst
  • Loading branch information
dcherian authored Oct 3, 2019
1 parent dd2b803 commit 86fb71d
Show file tree
Hide file tree
Showing 3 changed files with 71 additions and 0 deletions.
13 changes: 13 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
18 changes: 18 additions & 0 deletions xarray/core/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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

Expand Down
40 changes: 40 additions & 0 deletions xarray/tests/test_groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 86fb71d

Please sign in to comment.