TST: unit testing for *args/**kwargs in groupby. address GH #132

Author: wesm

pandas/core/groupby.py

@@ -863,9 +863,13 @@ def _aggregate_generic(self, func, *args, **kwargs):
             else:
                 raise e1
 
-        result = DataFrame(result)
-        if axis == 0:
-            result = result.T
+        if result:
+            if axis == 0:
+                result = DataFrame(result, index=obj.columns).T
+            else:
+                result = DataFrame(result, index=obj.index)
+        else:
+            result = DataFrame(result)
 
         return result
 

pandas/core/series.py

@@ -240,10 +240,12 @@ def __getitem__(self, key):
                 values = self.values
                 try:
                     return values[self.index.get_loc(key)]
-                except KeyError:
-                    if isinstance(key, (int, np.integer)):
+                except KeyError, e1:
+                    try:
                         return values[key]
-                    raise
+                    except Exception, _:
+                        pass
+                    raise e1
         except TypeError:
             pass
 

pandas/tests/test_groupby.py

@@ -753,6 +753,46 @@ def test_groupby_with_hier_columns(self):
         result = df.groupby(level=0).mean()
         self.assert_(result.columns.equals(sorted_columns))
 
+    def test_pass_args_kwargs(self):
+        from scipy.stats import scoreatpercentile
+
+        def f(x, q=None):
+            return scoreatpercentile(x, q)
+        g = lambda x: scoreatpercentile(x, 80)
+
+        # Series
+        ts_grouped = self.ts.groupby(lambda x: x.month)
+        agg_result = ts_grouped.agg(scoreatpercentile, 80)
+        apply_result = ts_grouped.apply(scoreatpercentile, 80)
+        trans_result = ts_grouped.transform(scoreatpercentile, 80)
+
+        agg_expected = ts_grouped.quantile(.8)
+        trans_expected = ts_grouped.transform(g)
+
+        assert_series_equal(apply_result, agg_expected)
+        assert_series_equal(agg_result, agg_expected)
+        assert_series_equal(trans_result, trans_expected)
+
+        agg_result = ts_grouped.agg(f, q=80)
+        apply_result = ts_grouped.apply(f, q=80)
+        trans_result = ts_grouped.transform(f, q=80)
+        assert_series_equal(agg_result, agg_expected)
+        assert_series_equal(apply_result, agg_expected)
+        assert_series_equal(trans_result, trans_expected)
+
+        # DataFrame
+        df_grouped = self.tsframe.groupby(lambda x: x.month)
+        agg_result = df_grouped.agg(scoreatpercentile, 80)
+        apply_result = df_grouped.apply(DataFrame.quantile, .8)
+        expected = df_grouped.quantile(.8)
+        assert_frame_equal(apply_result, expected)
+        assert_frame_equal(agg_result, expected)
+
+        agg_result = df_grouped.agg(f, q=80)
+        apply_result = df_grouped.apply(DataFrame.quantile, q=.8)
+        assert_frame_equal(agg_result, expected)
+        assert_frame_equal(apply_result, expected)
+
 class TestPanelGroupBy(unittest.TestCase):
 
     def setUp(self):