TYP: NDFrame.pipe, GroupBy.pipe etc. (#39093)

Author: topper-123

pandas/core/common.py

@@ -8,7 +8,18 @@
 import contextlib
 from functools import partial
 import inspect
-from typing import Any, Collection, Iterable, Iterator, List, Optional, Union, cast
+from typing import (
+    Any,
+    Callable,
+    Collection,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
 import warnings
 
 import numpy as np
@@ -405,7 +416,9 @@ def random_state(state=None):
         )
 
 
-def pipe(obj, func, *args, **kwargs):
+def pipe(
+    obj, func: Union[Callable[..., T], Tuple[Callable[..., T], str]], *args, **kwargs
+) -> T:
     """
     Apply a function ``func`` to object ``obj`` either by passing obj as the
     first argument to the function or, in the case that the func is a tuple,

pandas/core/generic.py

@@ -49,6 +49,7 @@
     NpDtype,
     Renamer,
     StorageOptions,
+    T,
     TimedeltaConvertibleTypes,
     TimestampConvertibleTypes,
     ValueKeyFunc,
@@ -5359,7 +5360,12 @@ def sample(
 
     @final
     @doc(klass=_shared_doc_kwargs["klass"])
-    def pipe(self, func, *args, **kwargs):
+    def pipe(
+        self,
+        func: Union[Callable[..., T], Tuple[Callable[..., T], str]],
+        *args,
+        **kwargs,
+    ) -> T:
         r"""
         Apply func(self, \*args, \*\*kwargs).
 

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 __future__ import annotations
 
 from contextlib import contextmanager
 import datetime
@@ -44,6 +45,7 @@ class providing the base-class of operations.
     FrameOrSeriesUnion,
     IndexLabel,
     Scalar,
+    T,
     final,
 )
 from pandas.compat.numpy import function as nv
@@ -475,7 +477,7 @@ def f(self):
 
 
 @contextmanager
-def group_selection_context(groupby: "BaseGroupBy") -> Iterator["BaseGroupBy"]:
+def group_selection_context(groupby: BaseGroupBy) -> Iterator[BaseGroupBy]:
     """
     Set / reset the group_selection_context.
     """
@@ -723,8 +725,8 @@ def _set_group_selection(self) -> None:
 
     @final
     def _set_result_index_ordered(
-        self, result: "OutputFrameOrSeries"
-    ) -> "OutputFrameOrSeries":
+        self, result: OutputFrameOrSeries
+    ) -> OutputFrameOrSeries:
         # set the result index on the passed values object and
         # return the new object, xref 8046
 
@@ -789,7 +791,12 @@ def __getattr__(self, attr: str):
         ),
     )
     @Appender(_pipe_template)
-    def pipe(self, func, *args, **kwargs):
+    def pipe(
+        self,
+        func: Union[Callable[..., T], Tuple[Callable[..., T], str]],
+        *args,
+        **kwargs,
+    ) -> T:
         return com.pipe(self, func, *args, **kwargs)
 
     plot = property(GroupByPlot)
@@ -3057,7 +3064,7 @@ def get_groupby(
     by: Optional[_KeysArgType] = None,
     axis: int = 0,
     level=None,
-    grouper: "Optional[ops.BaseGrouper]" = None,
+    grouper: Optional[ops.BaseGrouper] = None,
     exclusions=None,
     selection=None,
     as_index: bool = True,

pandas/core/resample.py

@@ -1,7 +1,9 @@
+from __future__ import annotations
+
 import copy
 from datetime import timedelta
 from textwrap import dedent
-from typing import Dict, Optional, Union, no_type_check
+from typing import Callable, Dict, Optional, Tuple, Union, no_type_check
 
 import numpy as np
 
@@ -14,7 +16,7 @@
     Timestamp,
     to_offset,
 )
-from pandas._typing import TimedeltaConvertibleTypes, TimestampConvertibleTypes
+from pandas._typing import T, TimedeltaConvertibleTypes, TimestampConvertibleTypes
 from pandas.compat.numpy import function as nv
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import Appender, Substitution, doc
@@ -231,7 +233,12 @@ def _assure_grouper(self):
     2012-08-04  1""",
     )
     @Appender(_pipe_template)
-    def pipe(self, func, *args, **kwargs):
+    def pipe(
+        self,
+        func: Union[Callable[..., T], Tuple[Callable[..., T], str]],
+        *args,
+        **kwargs,
+    ) -> T:
         return super().pipe(func, *args, **kwargs)
 
     _agg_see_also_doc = dedent(