updated test, removed last instance of not using context manager to f…

…ix funcs
pandas-dev · christopherzimmerman · Sep 19, 2019 · Sep 19, 2019 · Sep 19, 2019 · Sep 20, 2019
commit ca26b64bc5ee3b0e9dce9e7e820c3d1d48ebbba6
diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
@@ -606,56 +606,55 @@ def _make_wrapper(self, name):
             )
             raise AttributeError(msg)
 
-        self._set_group_selection()
-
-        # need to setup the selection
-        # as are not passed directly but in the grouper
-        f = getattr(self._selected_obj, name)
-        if not isinstance(f, types.MethodType):
-            return self.apply(lambda self: getattr(self, name))
-
-        f = getattr(type(self._selected_obj), name)
-
-        def wrapper(*args, **kwargs):
-            # a little trickery for aggregation functions that need an axis
-            # argument
-            kwargs_with_axis = kwargs.copy()
-            if "axis" not in kwargs_with_axis or kwargs_with_axis["axis"] is None:
-                kwargs_with_axis["axis"] = self.axis
-
-            def curried_with_axis(x):
-                return f(x, *args, **kwargs_with_axis)
-
-            def curried(x):
-                return f(x, *args, **kwargs)
-
-            # preserve the name so we can detect it when calling plot methods,
-            # to avoid duplicates
-            curried.__name__ = curried_with_axis.__name__ = name
-
-            # special case otherwise extra plots are created when catching the
-            # exception below
-            if name in base.plotting_methods:
-                return self.apply(curried)
+        with _group_selection_context(self):
+            # need to setup the selection
+            # as are not passed directly but in the grouper
+            f = getattr(self._selected_obj, name)
+            if not isinstance(f, types.MethodType):
+                return self.apply(lambda self: getattr(self, name))
+
+            f = getattr(type(self._selected_obj), name)
+
+            def wrapper(*args, **kwargs):
+                # a little trickery for aggregation functions that need an axis
+                # argument
+                kwargs_with_axis = kwargs.copy()
+                if "axis" not in kwargs_with_axis or kwargs_with_axis["axis"] is None:
+                    kwargs_with_axis["axis"] = self.axis
+
+                def curried_with_axis(x):
+                    return f(x, *args, **kwargs_with_axis)
+
+                def curried(x):
+                    return f(x, *args, **kwargs)
+
+                # preserve the name so we can detect it when calling plot methods,
+                # to avoid duplicates
+                curried.__name__ = curried_with_axis.__name__ = name
+
+                # special case otherwise extra plots are created when catching the
+                # exception below
+                if name in base.plotting_methods:
+                    return self.apply(curried)
 
-            try:
-                return self.apply(curried_with_axis)
-            except Exception:
                 try:
-                    return self.apply(curried)
+                    return self.apply(curried_with_axis)
                 except Exception:
-
-                    # related to : GH3688
-                    # try item-by-item
-                    # this can be called recursively, so need to raise
-                    # ValueError
-                    # if we don't have this method to indicated to aggregate to
-                    # mark this column as an error
                     try:
-                        return self._aggregate_item_by_item(name, *args, **kwargs)
-                    except AttributeError:
-                        # e.g. SparseArray has no flags attr
-                        raise ValueError
+                        return self.apply(curried)
+                    except Exception:
+
+                        # related to : GH3688
+                        # try item-by-item
+                        # this can be called recursively, so need to raise
+                        # ValueError
+                        # if we don't have this method to indicated to aggregate to
+                        # mark this column as an error
+                        try:
+                            return self._aggregate_item_by_item(name, *args, **kwargs)
+                        except AttributeError:
+                            # e.g. SparseArray has no flags attr
+                            raise ValueError
 
         return wrapper
 

diff --git a/pandas/tests/groupby/test_groupby.py b/pandas/tests/groupby/test_groupby.py
@@ -433,25 +433,15 @@ def test_frame_groupby_columns(tsframe):
         assert len(v.columns) == 2
 
 
-@pytest.mark.parametrize(
-    "func, args",
-    [
-        ("sum", []),
-        ("prod", []),
-        ("min", []),
-        ("max", []),
-        ("nth", [0]),
-        ("last", []),
-        ("first", []),
-    ],
-)
-def test_frame_groupby_avoids_mutate(func, args):
+def test_frame_groupby_avoids_mutate(reduction_func):
     # GH28523
+    func = reduction_func
     df = pd.DataFrame({"A": ["foo", "bar", "foo", "bar"], "B": [1, 2, 3, 4]})
     grouped = df.groupby("A")
 
     expected = grouped.apply(lambda x: x)
 
+    args = {"nth": [0], "quantile": [0.5]}.get(func, [])
     fn = getattr(grouped, func)
     fn(*args)