BUG: handle ndarrays/Index objects also in GroupBy.__getitem__, close #1257

Author: wesm

pandas/core/groupby.py

@@ -1656,7 +1656,7 @@ def __getitem__(self, key):
         if self._selection is not None:
             raise Exception('Column(s) %s already selected' % self._selection)
 
-        if isinstance(key, (list, tuple)) or not self.as_index:
+        if isinstance(key, (list, tuple, np.ndarray)) or not self.as_index:
             return DataFrameGroupBy(self.obj, self.grouper, selection=key,
                                     grouper=self.grouper,
                                     exclusions=self.exclusions,

pandas/tests/test_groupby.py

@@ -1724,10 +1724,13 @@ def test_getitem_list_of_columns(self):
 
         result = df.groupby('A')[['C', 'D']].mean()
         result2 = df.groupby('A')['C', 'D'].mean()
+        result3 = df.groupby('A')[df.columns[2:4]].mean()
+
         expected = df.ix[:, ['A', 'C', 'D']].groupby('A').mean()
 
         assert_frame_equal(result, expected)
         assert_frame_equal(result2, expected)
+        assert_frame_equal(result3, expected)
 
     def test_agg_multiple_functions_maintain_order(self):
         # GH #610