From 8dbef5b2cce01c13034580822e4851b04b457d69 Mon Sep 17 00:00:00 2001 From: Brock Date: Fri, 9 Apr 2021 09:35:35 -0700 Subject: [PATCH 1/2] TYP: implement typing.Manager2D --- pandas/_typing.py | 3 +++ pandas/core/groupby/generic.py | 16 +++++++++++----- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/pandas/_typing.py b/pandas/_typing.py index 36441e620286d..c8df3ea0d7f76 100644 --- a/pandas/_typing.py +++ b/pandas/_typing.py @@ -187,6 +187,9 @@ # internals Manager = Union["ArrayManager", "BlockManager", "SingleBlockManager"] SingleManager = Union["SingleArrayManager", "SingleBlockManager"] +Manager2D = Union["ArrayManager", "BlockManager"] +# TODO: Manager2d excludes SingleBlockManager, but does not exclude +# SingleArrayManager # indexing # PositionalIndexer -> valid 1D positional indexer, e.g. can pass diff --git a/pandas/core/groupby/generic.py b/pandas/core/groupby/generic.py index 5aebad84a0a30..5339546091f95 100644 --- a/pandas/core/groupby/generic.py +++ b/pandas/core/groupby/generic.py @@ -36,7 +36,7 @@ ArrayLike, FrameOrSeries, FrameOrSeriesUnion, - Manager, + Manager2D, ) from pandas.util._decorators import ( Appender, @@ -177,6 +177,9 @@ def pinner(cls): class SeriesGroupBy(GroupBy[Series]): _apply_allowlist = base.series_apply_allowlist + # Defined as a cache_readonly in SelectionMixin + _obj_with_exclusions: DataFrame + def _iterate_slices(self) -> Iterable[Series]: yield self._selected_obj @@ -927,6 +930,9 @@ def pct_change(self, periods=1, fill_method="pad", limit=None, freq=None): @pin_allowlisted_properties(DataFrame, base.dataframe_apply_allowlist) class DataFrameGroupBy(GroupBy[DataFrame]): + # Defined as a cache_readonly in SelectionMixin + _obj_with_exclusions: DataFrame + _apply_allowlist = base.dataframe_apply_allowlist _agg_examples_doc = dedent( @@ -1095,9 +1101,9 @@ def _cython_agg_general( def _cython_agg_manager( self, how: str, alt=None, numeric_only: bool = True, min_count: int = -1 - ) -> Manager: + ) -> Manager2D: - data: Manager = self._get_data_to_aggregate() + data: Manager2D = self._get_data_to_aggregate() if numeric_only: data = data.get_numeric_data(copy=False) @@ -1691,7 +1697,7 @@ def _wrap_frame_output(self, result, obj: DataFrame) -> DataFrame: else: return self.obj._constructor(result, index=obj.index, columns=result_index) - def _get_data_to_aggregate(self) -> Manager: + def _get_data_to_aggregate(self) -> Manager2D: obj = self._obj_with_exclusions if self.axis == 1: return obj.T._mgr @@ -1776,7 +1782,7 @@ def _wrap_transformed_output( return result - def _wrap_agged_manager(self, mgr: Manager) -> DataFrame: + def _wrap_agged_manager(self, mgr: Manager2D) -> DataFrame: if not self.as_index: index = np.arange(mgr.shape[1]) mgr.set_axis(1, ibase.Index(index), verify_integrity=False) From 332f395b2690f9aaa6d42b160697ba0b45e356a9 Mon Sep 17 00:00:00 2001 From: Brock Date: Fri, 9 Apr 2021 09:41:20 -0700 Subject: [PATCH 2/2] revert experiment --- pandas/core/groupby/generic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas/core/groupby/generic.py b/pandas/core/groupby/generic.py index 5339546091f95..4a2370d55aceb 100644 --- a/pandas/core/groupby/generic.py +++ b/pandas/core/groupby/generic.py @@ -178,7 +178,7 @@ class SeriesGroupBy(GroupBy[Series]): _apply_allowlist = base.series_apply_allowlist # Defined as a cache_readonly in SelectionMixin - _obj_with_exclusions: DataFrame + _obj_with_exclusions: Series def _iterate_slices(self) -> Iterable[Series]: yield self._selected_obj