ENH: column selection via attributes on DataFrameGroupBy, GH #882

Author: wesm

pandas/core/groupby.py

@@ -139,8 +139,12 @@ def _obj_with_exclusions(self):
         return self.obj
 
     def __getattr__(self, attr):
+        if attr in self.obj:
+            return self[attr]
+
         if hasattr(self.obj, attr) and attr != '_cache':
             return self._make_wrapper(attr)
+
         raise AttributeError("'%s' object has no attribute '%s'" %
                              (type(self).__name__, attr))
 
@@ -1644,3 +1648,27 @@ def numpy_groupby(data, labels, axis=0):
     group_sums = np.add.reduceat(ordered_data, groups_at, axis=axis)
 
     return group_sums
+
+from pandas.util import py3compat
+import sys
+
+def install_ipython_completers():  # pragma: no cover
+    """Register the DataFrame type with IPython's tab completion machinery, so
+    that it knows about accessing column names as attributes."""
+    from IPython.utils.generics import complete_object
+
+    @complete_object.when_type(DataFrameGroupBy)
+    def complete_dataframe(obj, prev_completions):
+        return prev_completions + [c for c in obj.obj.columns \
+                    if isinstance(c, basestring) and py3compat.isidentifier(c)]
+
+
+# Importing IPython brings in about 200 modules, so we want to avoid it unless
+# we're in IPython (when those modules are loaded anyway).
+if "IPython" in sys.modules:  # pragma: no cover
+    try:
+        install_ipython_completers()
+    except Exception:
+        pass
+
+

pandas/tests/test_groupby.py

@@ -1542,6 +1542,16 @@ def test_intercept_builtin_sum(self):
         assert_series_equal(result, expected)
         assert_series_equal(result2, expected)
 
+    def test_column_select_via_attr(self):
+        result = self.df.groupby('A').C.sum()
+        expected = self.df.groupby('A')['C'].sum()
+        assert_series_equal(result, expected)
+
+        self.df['mean'] = 1.5
+        result = self.df.groupby('A').mean()
+        expected = self.df.groupby('A').agg(np.mean)
+        assert_frame_equal(result, expected)
+
 def _check_groupby(df, result, keys, field, f=lambda x: x.sum()):
     tups = map(tuple, df[keys].values)
     tups = com._asarray_tuplesafe(tups)