Wip

Author: MarcoGorelli

pandas/core/base.py

@@ -214,11 +214,10 @@ def ndim(self) -> int:
 
     @cache_readonly
     def _obj_with_exclusions(self):
-        # breakpoint()
         if self._selection is not None and isinstance(self.obj, ABCDataFrame):
             return self.obj.reindex(columns=self._selection_list)
 
-        if len(self.exclusions) > 0 and isinstance(self.obj, ABCDataFrame):
+        if len(self.exclusions) > 0:
             return self.obj.drop(self.exclusions, axis=1)
         else:
             return self.obj

pandas/core/frame.py

@@ -7732,6 +7732,7 @@ def count(self, axis=0, level=None, numeric_only=False):
         Lewis     1
         Myla      1
         """
+        breakpoint()
         axis = self._get_axis_number(axis)
         if level is not None:
             return self._count_level(level, axis=axis, numeric_only=numeric_only)

pandas/core/groupby/groupby.py

@@ -6,6 +6,7 @@ class providing the base-class of operations.
 (defined in pandas.core.groupby.generic)
 expose these user-facing objects to provide specific functionality.
 """
+from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
 
 from contextlib import contextmanager
 import datetime
@@ -408,7 +409,6 @@ def __init__(
                 observed=observed,
                 mutated=self.mutated,
             )
-
         self.obj = obj
         self.axis = obj._get_axis_number(axis)
         self.grouper = grouper
@@ -508,13 +508,17 @@ def _get_index(self, name):
     @cache_readonly
     def _selected_obj(self):
         # Note: _selected_obj is always just `self.obj` for SeriesGroupBy
-
-        if self._selection is None or isinstance(self.obj, Series):
+        breakpoint()
+        if self.as_index:
+            obj = self._obj_with_exclusions
+        else:
+            obj = self.obj
+        if self._selection is None or isinstance(obj, Series):
             if self._group_selection is not None:
-                return self.obj[self._group_selection]
-            return self.obj
+                return obj[self._group_selection]
+            return obj
         else:
-            return self.obj[self._selection]
+            return obj[self._selection]
 
     def _reset_group_selection(self):
         """
@@ -704,7 +708,7 @@ def __iter__(self):
         )
     )
     def apply(self, func, *args, **kwargs):
-        # breakpoint()
+
         func = self._is_builtin_func(func)
 
         # this is needed so we don't try and wrap strings. If we could
@@ -732,7 +736,6 @@ def f(g):
         # ignore SettingWithCopy here in case the user mutates
         with option_context("mode.chained_assignment", None):
             try:
-                # breakpoint()
                 result = self._python_apply_general(f)
             except TypeError:
                 # gh-20949
@@ -749,11 +752,7 @@ def f(g):
         return result
 
     def _python_apply_general(self, f):
-        breakpoint()
-        if self.group_keys:
-            keys, values, mutated = self.grouper.apply(f, self._obj_with_exclusions, self.axis)
-        else:
-            keys, values, mutated = self.grouper.apply(f, self._selected_obj, self.axis)
+        keys, values, mutated = self.grouper.apply(f, self._selected_obj, self.axis)
 
         return self._wrap_applied_output(
             keys, values, not_indexed_same=mutated or self.mutated
@@ -1581,7 +1580,7 @@ def rolling(self, *args, **kwargs):
         Return a rolling grouper, providing rolling functionality per group.
         """
         from pandas.core.window import RollingGroupby
-        kwargs['exclusions'] = self.exclusions
+
         return RollingGroupby(self, *args, **kwargs)
 
     @Substitution(name="groupby")

pandas/core/groupby/ops.py

@@ -149,7 +149,6 @@ def _get_group_keys(self):
             return get_flattened_iterator(comp_ids, ngroups, self.levels, self.codes)
 
     def apply(self, f, data: FrameOrSeries, axis: int = 0):
-        # breakpoint()
         mutated = self.mutated
         splitter = self._get_splitter(data, axis=axis)
         group_keys = self._get_group_keys()
@@ -170,7 +169,6 @@ def apply(self, f, data: FrameOrSeries, axis: int = 0):
             and not sdata.index._has_complex_internals
         ):
             try:
-                # breakpoint()
                 result_values, mutated = splitter.fast_apply(f, sdata, group_keys)
 
             except libreduction.InvalidApply as err:
@@ -929,7 +927,6 @@ def _chop(self, sdata: Series, slice_obj: slice) -> Series:
 class FrameSplitter(DataSplitter):
     def fast_apply(self, f, sdata: FrameOrSeries, names):
         # must return keys::list, values::list, mutated::bool
-        # breakpoint()
         starts, ends = lib.generate_slices(self.slabels, self.ngroups)
         return libreduction.apply_frame_axis0(sdata, f, names, starts, ends)
 

pandas/core/window/common.py

@@ -52,7 +52,7 @@ def __init__(self, obj, *args, **kwargs):
         kwargs.pop("parent", None)
         groupby = kwargs.pop("groupby", None)
         if groupby is None:
-            groupby, obj = obj, obj._obj_with_exclusions
+            groupby, obj = obj, obj.obj
         self._groupby = groupby
         self._groupby.mutated = True
         self._groupby.grouper.mutated = True
@@ -82,13 +82,12 @@ def _apply(
         # TODO: can we de-duplicate with _dispatch?
         def f(x, name=name, *args):
             x = self._shallow_copy(x)
-            # x.obj = x._obj_with_exclusions
 
             if isinstance(name, str):
                 return getattr(x, name)(*args, **kwargs)
 
             return x.apply(name, *args, **kwargs)
-        # breakpoint()
+
         return self._groupby.apply(f)
 
 

pandas/core/window/rolling.py

@@ -93,7 +93,6 @@ def __init__(
         self.axis = obj._get_axis_number(axis) if axis is not None else None
         self.validate()
         self._numba_func_cache: Dict[Optional[str], Callable] = dict()
-        self.exclusions = kwargs.get('exclusions', set())
 
     @property
     def _constructor(self):
@@ -1174,7 +1173,6 @@ class _Rolling_and_Expanding(_Rolling):
     )
 
     def count(self):
-
         blocks, obj = self._create_blocks()
         results = []
         for b in blocks:
@@ -1188,7 +1186,6 @@ def count(self):
                 closed=self.closed,
             ).sum()
             results.append(result)
-
         return self._wrap_results(results, blocks, obj)
 
     _shared_docs["apply"] = dedent(
@@ -1954,7 +1951,6 @@ def aggregate(self, func, *args, **kwargs):
     @Substitution(name="rolling")
     @Appender(_shared_docs["count"])
     def count(self):
-
         # different impl for freq counting
         if self.is_freq_type:
             window_func = self._get_roll_func("roll_count")
@@ -1972,8 +1968,7 @@ def apply(
         engine_kwargs=None,
         args=None,
         kwargs=None,
-    ):  
-        # breakpoint()
+    ):
         return super().apply(
             func,
             raw=raw,

pandas/tests/groupby/test_apply.py

@@ -370,7 +370,6 @@ def test_apply_chunk_view():
 
     result = df.groupby("key", group_keys=False).apply(lambda x: x[:2])
     expected = df.take([0, 1, 3, 4, 6, 7])
-    breakpoint()
     tm.assert_frame_equal(result, expected)
 
 
@@ -406,7 +405,6 @@ def f(group):
 
     expected = df.copy()
     expected["v2"] = np.tile([0.0, 0.5, 1], 2)
-    breakpoint()
     tm.assert_frame_equal(result, expected)
 
 
@@ -828,6 +826,7 @@ def test_apply_index_has_complex_internals(index):
     # GH 31248
     df = DataFrame({"group": [1, 1, 2], "value": [0, 1, 0]}, index=index)
     result = df.groupby("group").apply(lambda x: x)
+    breakpoint()
     tm.assert_frame_equal(result, df)
 
 

pandas/tests/window/test_grouper.py

@@ -30,7 +30,6 @@ def test_getitem(self):
         expected = g_mutated.B.apply(lambda x: x.rolling(2).mean())
 
         result = g.rolling(2).mean().B
-        # breakpoint()
         tm.assert_series_equal(result, expected)
 
         result = g.rolling(2).B.mean()
@@ -62,9 +61,7 @@ def test_rolling(self):
 
         for f in ["sum", "mean", "min", "max", "count", "kurt", "skew"]:
             result = getattr(r, f)()
-            # breakpoint()
             expected = g.apply(lambda x: getattr(x.rolling(4), f)())
-            # breakpoint()
             tm.assert_frame_equal(result, expected)
 
         for f in ["std", "var"]: