@@ -736,13 +736,12 @@ def pipe(self, func, *args, **kwargs):
736
736
def _make_wrapper (self , name ):
737
737
assert name in self ._apply_allowlist
738
738
739
- self ._set_group_selection ()
740
-
741
- # need to setup the selection
742
- # as are not passed directly but in the grouper
743
- f = getattr (self ._obj_with_exclusions , name )
744
- if not isinstance (f , types .MethodType ):
745
- return self .apply (lambda self : getattr (self , name ))
739
+ with _group_selection_context (self ):
740
+ # need to setup the selection
741
+ # as are not passed directly but in the grouper
742
+ f = getattr (self ._obj_with_exclusions , name )
743
+ if not isinstance (f , types .MethodType ):
744
+ return self .apply (lambda self : getattr (self , name ))
746
745
747
746
f = getattr (type (self ._obj_with_exclusions ), name )
748
747
sig = inspect .signature (f )
@@ -992,28 +991,28 @@ def _agg_general(
992
991
alias : str ,
993
992
npfunc : Callable ,
994
993
):
995
- self ._set_group_selection ()
996
-
997
- # try a cython aggregation if we can
998
- try :
999
- return self ._cython_agg_general (
1000
- how = alias , alt = npfunc , numeric_only = numeric_only , min_count = min_count ,
1001
- )
1002
- except DataError :
1003
- pass
1004
- except NotImplementedError as err :
1005
- if "function is not implemented for this dtype" in str (
1006
- err
1007
- ) or "category dtype not supported" in str (err ):
1008
- # raised in _get_cython_function, in some cases can
1009
- # be trimmed by implementing cython funcs for more dtypes
994
+ with _group_selection_context (self ):
995
+ # try a cython aggregation if we can
996
+ try :
997
+ return self ._cython_agg_general (
998
+ how = alias ,
999
+ alt = npfunc ,
1000
+ numeric_only = numeric_only ,
1001
+ min_count = min_count ,
1002
+ )
1003
+ except DataError :
1010
1004
pass
1011
- else :
1012
- raise
1013
-
1014
- # apply a non-cython aggregation
1015
- result = self .aggregate (lambda x : npfunc (x , axis = self .axis ))
1016
- return result
1005
+ except NotImplementedError as err :
1006
+ if "function is not implemented for this dtype" in str (
1007
+ err
1008
+ ) or "category dtype not supported" in str (err ):
1009
+ # raised in _get_cython_function, in some cases can
1010
+ # be trimmed by implementing cython funcs for more dtypes
1011
+ pass
1012
+
1013
+ # apply a non-cython aggregation
1014
+ result = self .aggregate (lambda x : npfunc (x , axis = self .axis ))
1015
+ return result
1017
1016
1018
1017
def _cython_agg_general (
1019
1018
self , how : str , alt = None , numeric_only : bool = True , min_count : int = - 1
@@ -1940,29 +1939,31 @@ def nth(self, n: Union[int, List[int]], dropna: Optional[str] = None) -> DataFra
1940
1939
nth_values = list (set (n ))
1941
1940
1942
1941
nth_array = np .array (nth_values , dtype = np .intp )
1943
- self . _set_group_selection ()
1942
+ with _group_selection_context ( self ):
1944
1943
1945
- mask_left = np .in1d (self ._cumcount_array (), nth_array )
1946
- mask_right = np .in1d (self ._cumcount_array (ascending = False ) + 1 , - nth_array )
1947
- mask = mask_left | mask_right
1944
+ mask_left = np .in1d (self ._cumcount_array (), nth_array )
1945
+ mask_right = np .in1d (
1946
+ self ._cumcount_array (ascending = False ) + 1 , - nth_array
1947
+ )
1948
+ mask = mask_left | mask_right
1948
1949
1949
- ids , _ , _ = self .grouper .group_info
1950
+ ids , _ , _ = self .grouper .group_info
1950
1951
1951
- # Drop NA values in grouping
1952
- mask = mask & (ids != - 1 )
1952
+ # Drop NA values in grouping
1953
+ mask = mask & (ids != - 1 )
1953
1954
1954
- out = self ._selected_obj [mask ]
1955
- if not self .as_index :
1956
- return out
1955
+ out = self ._selected_obj [mask ]
1956
+ if not self .as_index :
1957
+ return out
1957
1958
1958
- result_index = self .grouper .result_index
1959
- out .index = result_index [ids [mask ]]
1959
+ result_index = self .grouper .result_index
1960
+ out .index = result_index [ids [mask ]]
1960
1961
1961
- if not self .observed and isinstance (result_index , CategoricalIndex ):
1962
- out = out .reindex (result_index )
1962
+ if not self .observed and isinstance (result_index , CategoricalIndex ):
1963
+ out = out .reindex (result_index )
1963
1964
1964
- out = self ._reindex_output (out )
1965
- return out .sort_index () if self .sort else out
1965
+ out = self ._reindex_output (out )
1966
+ return out .sort_index () if self .sort else out
1966
1967
1967
1968
# dropna is truthy
1968
1969
if isinstance (n , valid_containers ):
0 commit comments