BUG: allow plot, boxplot, hist and completion on GroupBy objects

Author: cpcloud

pandas/core/groupby.py

@@ -1,4 +1,5 @@
 import types
+from functools import wraps
 import numpy as np
 
 from pandas.compat import(
@@ -45,14 +46,19 @@
 """
 
 
+# special case to prevent duplicate plots when catching exceptions when
+# forwarding methods from NDFrames
+_plotting_methods = frozenset(['plot', 'boxplot', 'hist'])
+
 _apply_whitelist = frozenset(['last', 'first',
                               'mean', 'sum', 'min', 'max',
                               'head', 'tail',
                               'cumsum', 'cumprod', 'cummin', 'cummax',
                               'resample',
                               'describe',
                               'rank', 'quantile', 'count',
-                              'fillna', 'dtype'])
+                              'fillna', 'dtype']) | _plotting_methods
+
 
 
 class GroupByError(Exception):
@@ -180,7 +186,6 @@ class GroupBy(PandasObject):
     len(grouped) : int
         Number of groups
     """
-
     def __init__(self, obj, keys=None, axis=0, level=None,
                  grouper=None, exclusions=None, selection=None, as_index=True,
                  sort=True, group_keys=True, squeeze=False):
@@ -244,6 +249,9 @@ def _selection_list(self):
             return [self._selection]
         return self._selection
 
+    def _local_dir(self):
+        return sorted(set(self.obj._local_dir() + list(_apply_whitelist)))
+
     def __getattr__(self, attr):
         if attr in self.obj:
             return self[attr]
@@ -281,9 +289,16 @@ def wrapper(*args, **kwargs):
 
             def curried_with_axis(x):
                 return f(x, *args, **kwargs_with_axis)
+            curried_with_axis.__name__ = name
 
             def curried(x):
                 return f(x, *args, **kwargs)
+            curried.__name__ = name
+
+            # special case otherwise extra plots are created when catching the
+            # exception below
+            if name in _plotting_methods:
+                return self.apply(curried)
 
             try:
                 return self.apply(curried_with_axis)
@@ -348,7 +363,11 @@ def apply(self, func, *args, **kwargs):
         applied : type depending on grouped object and function
         """
         func = _intercept_function(func)
-        f = lambda g: func(g, *args, **kwargs)
+
+        @wraps(func)
+        def f(g):
+            return func(g, *args, **kwargs)
+
         return self._python_apply_general(f)
 
     def _python_apply_general(self, f):
@@ -598,7 +617,7 @@ def __iter__(self):
     def nkeys(self):
         return len(self.groupings)
 
-    def get_iterator(self, data, axis=0, keep_internal=True):
+    def get_iterator(self, data, axis=0):
         """
         Groupby iterator
 
@@ -607,16 +626,14 @@ def get_iterator(self, data, axis=0, keep_internal=True):
         Generator yielding sequence of (name, subsetted object)
         for each group
         """
-        splitter = self._get_splitter(data, axis=axis,
-                                      keep_internal=keep_internal)
+        splitter = self._get_splitter(data, axis=axis)
         keys = self._get_group_keys()
         for key, (i, group) in zip(keys, splitter):
             yield key, group
 
-    def _get_splitter(self, data, axis=0, keep_internal=True):
+    def _get_splitter(self, data, axis=0):
         comp_ids, _, ngroups = self.group_info
-        return get_splitter(data, comp_ids, ngroups, axis=axis,
-                            keep_internal=keep_internal)
+        return get_splitter(data, comp_ids, ngroups, axis=axis)
 
     def _get_group_keys(self):
         if len(self.groupings) == 1:
@@ -627,19 +644,19 @@ def _get_group_keys(self):
             mapper = _KeyMapper(comp_ids, ngroups, self.labels, self.levels)
             return [mapper.get_key(i) for i in range(ngroups)]
 
-    def apply(self, f, data, axis=0, keep_internal=False):
+    def apply(self, f, data, axis=0):
         mutated = False
-        splitter = self._get_splitter(data, axis=axis,
-                                      keep_internal=keep_internal)
+        splitter = self._get_splitter(data, axis=axis)
         group_keys = self._get_group_keys()
 
         # oh boy
-        if hasattr(splitter, 'fast_apply') and axis == 0:
+        if (f.__name__ not in _plotting_methods and
+            hasattr(splitter, 'fast_apply') and axis == 0):
             try:
                 values, mutated = splitter.fast_apply(f, group_keys)
                 return group_keys, values, mutated
-            except (Exception) as detail:
-                # we detect a mutatation of some kind
+            except Exception:
+                # we detect a mutation of some kind
                 # so take slow path
                 pass
 
@@ -1043,7 +1060,7 @@ def get_iterator(self, data, axis=0):
                 inds = lrange(start, n)
                 yield self.binlabels[-1], data.take(inds, axis=axis)
 
-    def apply(self, f, data, axis=0, keep_internal=False):
+    def apply(self, f, data, axis=0):
         result_keys = []
         result_values = []
         mutated = False
@@ -1617,6 +1634,7 @@ def filter(self, func, dropna=True, *args, **kwargs):
         else:
             return filtered.reindex(self.obj.index) # Fill with NaNs.
 
+
 class NDFrameGroupBy(GroupBy):
 
     def _iterate_slices(self):
@@ -2268,6 +2286,7 @@ def ohlc(self):
         """
         return self._apply_to_column_groupbys(lambda x: x._cython_agg_general('ohlc'))
 
+
 from pandas.tools.plotting import boxplot_frame_groupby
 DataFrameGroupBy.boxplot = boxplot_frame_groupby
 
@@ -2364,7 +2383,7 @@ class NDArrayGroupBy(GroupBy):
 
 class DataSplitter(object):
 
-    def __init__(self, data, labels, ngroups, axis=0, keep_internal=False):
+    def __init__(self, data, labels, ngroups, axis=0):
         self.data = data
         self.labels = com._ensure_int64(labels)
         self.ngroups = ngroups
@@ -2419,10 +2438,8 @@ def _chop(self, sdata, slice_obj):
 
 
 class FrameSplitter(DataSplitter):
-
-    def __init__(self, data, labels, ngroups, axis=0, keep_internal=False):
-        DataSplitter.__init__(self, data, labels, ngroups, axis=axis,
-                              keep_internal=keep_internal)
+    def __init__(self, data, labels, ngroups, axis=0):
+        super(FrameSplitter, self).__init__(data, labels, ngroups, axis=axis)
 
     def fast_apply(self, f, names):
         # must return keys::list, values::list, mutated::bool
@@ -2445,10 +2462,8 @@ def _chop(self, sdata, slice_obj):
 
 
 class NDFrameSplitter(DataSplitter):
-
-    def __init__(self, data, labels, ngroups, axis=0, keep_internal=False):
-        DataSplitter.__init__(self, data, labels, ngroups, axis=axis,
-                              keep_internal=keep_internal)
+    def __init__(self, data, labels, ngroups, axis=0):
+        super(NDFrameSplitter, self).__init__(data, labels, ngroups, axis=axis)
 
         self.factory = data._constructor
 

pandas/tests/test_groupby.py

@@ -2728,6 +2728,78 @@ def test_groupby_whitelist(self):
                 with tm.assertRaisesRegexp(AttributeError, msg):
                     getattr(gb, bl)
 
+    def test_series_groupby_plotting_nominally_works(self):
+        try:
+            import matplotlib.pyplot as plt
+        except ImportError:
+            raise nose.SkipTest("matplotlib not installed")
+        n = 10
+        weight = Series(np.random.normal(166, 20, size=n))
+        height = Series(np.random.normal(60, 10, size=n))
+        gender = tm.choice(['male', 'female'], size=n)
+
+        weight.groupby(gender).plot()
+        tm.close()
+        height.groupby(gender).hist()
+        tm.close()
+
+    def test_frame_groupby_plot_boxplot(self):
+        try:
+            import matplotlib.pyplot as plt
+        except ImportError:
+            raise nose.SkipTest("matplotlib not installed")
+        tm.close()
+
+        n = 10
+        weight = Series(np.random.normal(166, 20, size=n))
+        height = Series(np.random.normal(60, 10, size=n))
+        gender = tm.choice(['male', 'female'], size=n)
+        df = DataFrame({'height': height, 'weight': weight, 'gender': gender})
+        gb = df.groupby('gender')
+
+        res = gb.plot()
+        self.assertEqual(len(plt.get_fignums()), 2)
+        self.assertEqual(len(res), 2)
+        tm.close()
+
+        res = gb.boxplot()
+        self.assertEqual(len(plt.get_fignums()), 1)
+        self.assertEqual(len(res), 2)
+        tm.close()
+
+        with tm.assertRaises(TypeError, '.*str.+float'):
+            gb.hist()
+
+    def test_frame_groupby_hist(self):
+        try:
+            import matplotlib.pyplot as plt
+        except ImportError:
+            raise nose.SkipTest("matplotlib not installed")
+        tm.close()
+
+        n = 10
+        weight = Series(np.random.normal(166, 20, size=n))
+        height = Series(np.random.normal(60, 10, size=n))
+        gender_int = tm.choice([0, 1], size=n)
+        df_int = DataFrame({'height': height, 'weight': weight,
+                            'gender': gender_int})
+        gb = df_int.groupby('gender')
+        axes = gb.hist()
+        self.assertEqual(len(axes), 2)
+        self.assertEqual(len(plt.get_fignums()), 2)
+        tm.close()
+
+    def test_tab_completion(self):
+        grp = self.mframe.groupby(level='second')
+        results = set([v for v in grp.__dir__() if not v.startswith('_')])
+        expected = set(['A','B','C',
+            'agg','aggregate','apply','boxplot','filter','first','get_group',
+            'groups','hist','indices','last','max','mean','median',
+            'min','name','ngroups','nth','ohlc','plot', 'prod',
+            'size','std','sum','transform','var', 'count', 'head', 'describe',
+            'cummax', 'dtype', 'quantile', 'rank',
+            'cumprod', 'tail', 'resample', 'cummin', 'fillna', 'cumsum'])
+        self.assertEqual(results, expected)
 
 def assert_fp_equal(a, b):
     assert (np.abs(a - b) < 1e-12).all()