diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 79e941f262931..c82d8a25fedba 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -14,7 +14,18 @@ import itertools import sys from textwrap import dedent -from typing import FrozenSet, List, Optional, Sequence, Set, Tuple, Type, Union +from typing import ( + FrozenSet, + Hashable, + Iterable, + List, + Optional, + Sequence, + Set, + Tuple, + Type, + Union, +) import warnings import numpy as np @@ -861,7 +872,7 @@ def style(self): """ @Appender(_shared_docs["items"]) - def items(self): + def items(self) -> Iterable[Tuple[Hashable, Series]]: if self.columns.is_unique and hasattr(self, "_item_cache"): for k in self.columns: yield k, self._get_item_cache(k) diff --git a/pandas/core/groupby/generic.py b/pandas/core/groupby/generic.py index 068d5e5275f0d..4125d6a918b26 100644 --- a/pandas/core/groupby/generic.py +++ b/pandas/core/groupby/generic.py @@ -11,7 +11,17 @@ from functools import partial from textwrap import dedent import typing -from typing import Any, Callable, FrozenSet, Sequence, Type, Union +from typing import ( + Any, + Callable, + FrozenSet, + Hashable, + Iterable, + Sequence, + Tuple, + Type, + Union, +) import warnings import numpy as np @@ -132,7 +142,7 @@ def pinner(cls): class SeriesGroupBy(GroupBy): _apply_whitelist = base.series_apply_whitelist - def _iterate_slices(self): + def _iterate_slices(self) -> Iterable[Tuple[Hashable, Series]]: yield self._selection_name, self._selected_obj @property @@ -898,22 +908,20 @@ def aggregate(self, func=None, *args, **kwargs): agg = aggregate - def _iterate_slices(self): - if self.axis == 0: - # kludge - if self._selection is None: - slice_axis = self.obj.columns - else: - slice_axis = self._selection_list - slicer = lambda x: self.obj[x] + def _iterate_slices(self) -> Iterable[Tuple[Hashable, Series]]: + obj = self._selected_obj + if self.axis == 1: + obj = obj.T + + if isinstance(obj, Series) and obj.name not in self.exclusions: + # Occurs when doing DataFrameGroupBy(...)["X"] + yield obj.name, obj else: - slice_axis = self.obj.index - slicer = self.obj.xs + for label, values in obj.items(): + if label in self.exclusions: + continue - for val in slice_axis: - if val in self.exclusions: - continue - yield val, slicer(val) + yield label, values def _cython_agg_general(self, how, alt=None, numeric_only=True, min_count=-1): new_items, new_blocks = self._cython_agg_blocks( diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py index cc297629a7004..4ea02d59597f1 100644 --- a/pandas/core/groupby/groupby.py +++ b/pandas/core/groupby/groupby.py @@ -14,7 +14,7 @@ class providing the base-class of operations. import inspect import re import types -from typing import FrozenSet, List, Optional, Tuple, Type, Union +from typing import FrozenSet, Hashable, Iterable, List, Optional, Tuple, Type, Union import numpy as np @@ -758,7 +758,7 @@ def _python_apply_general(self, f): keys, values, not_indexed_same=mutated or self.mutated ) - def _iterate_slices(self): + def _iterate_slices(self) -> Iterable[Tuple[Hashable, Series]]: raise AbstractMethodError(self) def transform(self, func, *args, **kwargs):