TYP: misc return types (#57285)

Author: twoertwein

pandas/_typing.py

@@ -46,7 +46,12 @@
 
     from pandas.core.dtypes.dtypes import ExtensionDtype
 
-    from pandas import Interval
+    from pandas import (
+        DatetimeIndex,
+        Interval,
+        PeriodIndex,
+        TimedeltaIndex,
+    )
     from pandas.arrays import (
         DatetimeArray,
         TimedeltaArray,
@@ -190,6 +195,7 @@ def __reversed__(self) -> Iterator[_T_co]:
 NDFrameT = TypeVar("NDFrameT", bound="NDFrame")
 
 IndexT = TypeVar("IndexT", bound="Index")
+FreqIndexT = TypeVar("FreqIndexT", "DatetimeIndex", "PeriodIndex", "TimedeltaIndex")
 NumpyIndexT = TypeVar("NumpyIndexT", np.ndarray, "Index")
 
 AxisInt = int

pandas/core/_numba/extensions.py

@@ -12,6 +12,7 @@
 
 from contextlib import contextmanager
 import operator
+from typing import TYPE_CHECKING
 
 import numba
 from numba import types
@@ -40,6 +41,9 @@
 from pandas.core.internals import SingleBlockManager
 from pandas.core.series import Series
 
+if TYPE_CHECKING:
+    from pandas._typing import Self
+
 
 # Helper function to hack around fact that Index casts numpy string dtype to object
 #
@@ -84,7 +88,7 @@ def key(self):
     def as_array(self):
         return types.Array(self.dtype, 1, self.layout)
 
-    def copy(self, dtype=None, ndim: int = 1, layout=None):
+    def copy(self, dtype=None, ndim: int = 1, layout=None) -> Self:
         assert ndim == 1
         if dtype is None:
             dtype = self.dtype
@@ -114,7 +118,7 @@ def key(self):
     def as_array(self):
         return self.values
 
-    def copy(self, dtype=None, ndim: int = 1, layout: str = "C"):
+    def copy(self, dtype=None, ndim: int = 1, layout: str = "C") -> Self:
         assert ndim == 1
         assert layout == "C"
         if dtype is None:
@@ -123,7 +127,7 @@ def copy(self, dtype=None, ndim: int = 1, layout: str = "C"):
 
 
 @typeof_impl.register(Index)
-def typeof_index(val, c):
+def typeof_index(val, c) -> IndexType:
     """
     This will assume that only strings are in object dtype
     index.
@@ -136,7 +140,7 @@ def typeof_index(val, c):
 
 
 @typeof_impl.register(Series)
-def typeof_series(val, c):
+def typeof_series(val, c) -> SeriesType:
     index = typeof_impl(val.index, c)
     arrty = typeof_impl(val.values, c)
     namety = typeof_impl(val.name, c)
@@ -532,7 +536,7 @@ def key(self):
 
 
 @typeof_impl.register(_iLocIndexer)
-def typeof_iloc(val, c):
+def typeof_iloc(val, c) -> IlocType:
     objtype = typeof_impl(val.obj, c)
     return IlocType(objtype)
 

pandas/core/arrays/_mixins.py

@@ -250,7 +250,7 @@ def searchsorted(
         return self._ndarray.searchsorted(npvalue, side=side, sorter=sorter)
 
     @doc(ExtensionArray.shift)
-    def shift(self, periods: int = 1, fill_value=None):
+    def shift(self, periods: int = 1, fill_value=None) -> Self:
         # NB: shift is always along axis=0
         axis = 0
         fill_value = self._validate_scalar(fill_value)

pandas/core/arrays/arrow/array.py

@@ -204,7 +204,7 @@ def floordiv_compat(
     from pandas.core.arrays.timedeltas import TimedeltaArray
 
 
-def get_unit_from_pa_dtype(pa_dtype):
+def get_unit_from_pa_dtype(pa_dtype) -> str:
     # https://github.com/pandas-dev/pandas/pull/50998#discussion_r1100344804
     if pa_version_under11p0:
         unit = str(pa_dtype).split("[", 1)[-1][:-1]
@@ -1966,7 +1966,7 @@ def _rank(
         na_option: str = "keep",
         ascending: bool = True,
         pct: bool = False,
-    ):
+    ) -> Self:
         """
         See Series.rank.__doc__.
         """

pandas/core/dtypes/dtypes.py

@@ -2337,7 +2337,7 @@ def _get_common_dtype(self, dtypes: list[DtypeObj]) -> DtypeObj | None:
         except NotImplementedError:
             return None
 
-    def __from_arrow__(self, array: pa.Array | pa.ChunkedArray):
+    def __from_arrow__(self, array: pa.Array | pa.ChunkedArray) -> ArrowExtensionArray:
         """
         Construct IntegerArray/FloatingArray from pyarrow Array/ChunkedArray.
         """

pandas/core/dtypes/generic.py

@@ -33,7 +33,7 @@
 
 # define abstract base classes to enable isinstance type checking on our
 # objects
-def create_pandas_abc_type(name, attr, comp):
+def create_pandas_abc_type(name, attr, comp) -> type:
     def _check(inst) -> bool:
         return getattr(inst, attr, "_typ") in comp
 

pandas/core/groupby/generic.py

@@ -1788,7 +1788,7 @@ def _choose_path(self, fast_path: Callable, slow_path: Callable, group: DataFram
 
         return path, res
 
-    def filter(self, func, dropna: bool = True, *args, **kwargs):
+    def filter(self, func, dropna: bool = True, *args, **kwargs) -> DataFrame:
         """
         Filter elements from groups that don't satisfy a criterion.
 

pandas/core/indexing.py

@@ -1750,7 +1750,7 @@ def _get_slice_axis(self, slice_obj: slice, axis: AxisInt):
         labels._validate_positional_slice(slice_obj)
         return self.obj._slice(slice_obj, axis=axis)
 
-    def _convert_to_indexer(self, key, axis: AxisInt):
+    def _convert_to_indexer(self, key: T, axis: AxisInt) -> T:
         """
         Much simpler as we only have to deal with our valid types.
         """

pandas/core/interchange/from_dataframe.py

@@ -75,7 +75,7 @@ def from_dataframe(df, allow_copy: bool = True) -> pd.DataFrame:
     )
 
 
-def _from_dataframe(df: DataFrameXchg, allow_copy: bool = True):
+def _from_dataframe(df: DataFrameXchg, allow_copy: bool = True) -> pd.DataFrame:
     """
     Build a ``pd.DataFrame`` from the DataFrame interchange object.
 

pandas/core/internals/blocks.py

@@ -2537,7 +2537,7 @@ def get_block_type(dtype: DtypeObj) -> type[Block]:
 
 def new_block_2d(
     values: ArrayLike, placement: BlockPlacement, refs: BlockValuesRefs | None = None
-):
+) -> Block:
     # new_block specialized to case with
     #  ndim=2
     #  isinstance(placement, BlockPlacement)