From ed8d2263110c6b7a393eb6ae423b004925856118 Mon Sep 17 00:00:00 2001
From: dcherian <deepak@cherian.net>
Date: Sat, 6 Feb 2021 09:07:23 -0700
Subject: [PATCH] Implement GroupBy.__getitem__

---
 doc/groupby.rst              |  6 ++++++
 doc/whats-new.rst            |  4 ++++
 xarray/core/groupby.py       |  9 +++++++++
 xarray/tests/test_groupby.py | 13 +++++++++++++
 4 files changed, 32 insertions(+)

diff --git a/doc/groupby.rst b/doc/groupby.rst
index d0c0b1849f9..804799aebd1 100644
--- a/doc/groupby.rst
+++ b/doc/groupby.rst
@@ -63,6 +63,12 @@ You can also iterate over groups in ``(label, group)`` pairs:
 
     list(ds.groupby("letters"))
 
+You can index out a particular group:
+
+.. ipython:: python
+
+    ds.groupby("letters")["b"]
+
 Just like in pandas, creating a GroupBy object is cheap: it does not actually
 split the data until you access particular values.
 
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 471e91a8512..158ddf3350a 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -62,6 +62,10 @@ New Features
 - :py:meth:`DataArray.swap_dims` & :py:meth:`Dataset.swap_dims` now accept dims
   in the form of kwargs as well as a dict, like most similar methods.
   By `Maximilian Roos <https://github.com/max-sixty>`_.
+- Implement ``__getitem__`` for both :py:class:`~core.groupby.DatasetGroupBy` and
+  :py:class:`~core.groupby.DataArrayGroupBy`, inspired by pandas'
+  :py:meth:`~pandas.core.groupby.GroupBy.get_group`.
+  By `Deepak Cherian <https://github.com/dcherian>`_.
 
 Bug fixes
 ~~~~~~~~~
diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py
index e1e5a0fabe8..bfcbd4b62ee 100644
--- a/xarray/core/groupby.py
+++ b/xarray/core/groupby.py
@@ -415,11 +415,20 @@ def dims(self):
 
     @property
     def groups(self):
+        """
+        Mapping from group labels to indices. The indices can be used to index the underlying object.
+        """
         # provided to mimic pandas.groupby
         if self._groups is None:
             self._groups = dict(zip(self._unique_coord.values, self._group_indices))
         return self._groups
 
+    def __getitem__(self, key):
+        """
+        Get DataArray or Dataset corresponding to a particular group label.
+        """
+        return self._obj.isel({self._group_dim: self.groups[key]})
+
     def __len__(self):
         return self._unique_coord.size
 
diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py
index 85f729a9f7a..1366dd417f7 100644
--- a/xarray/tests/test_groupby.py
+++ b/xarray/tests/test_groupby.py
@@ -549,4 +549,17 @@ def test_groupby_none_group_name():
     assert "group" in mean.dims
 
 
+def test_groupby_sel(dataset):
+
+    assert_identical(dataset.sel(x="a"), dataset.groupby("x")["a"])
+    assert_identical(dataset.sel(z=1), dataset.groupby("z")[1])
+
+    assert_identical(dataset.foo.sel(x="a"), dataset.foo.groupby("x")["a"])
+    assert_identical(dataset.foo.sel(z=1), dataset.foo.groupby("z")[1])
+
+    actual = dataset.groupby("boo")["f"].unstack().transpose("x", "y", "z")
+    expected = dataset.sel(y=[1], z=[1, 2]).transpose("x", "y", "z")
+    assert_identical(expected, actual)
+
+
 # TODO: move other groupby tests from test_dataset and test_dataarray over here