CLN: Consistent and Annotated Return Type of _iterate_slices (pandas-dev#28958)

Author: WillAyd

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)

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(

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
 
@@ -751,7 +751,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):