Typing improvements for Index

Author: AndreyKolomiets

pandas/core/indexes/base.py

@@ -11,6 +11,8 @@
     ClassVar,
     Literal,
     NoReturn,
+    Optional,
+    Union,
     cast,
     final,
     overload,
@@ -46,6 +48,7 @@
     Axes,
     Axis,
     DropKeep,
+    Dtype,
     DtypeObj,
     F,
     IgnoreRaise,
@@ -155,6 +158,8 @@
     ExtensionArray,
     TimedeltaArray,
 )
+from pandas.core.arrays.floating import FloatingDtype
+from pandas.core.arrays.integer import IntegerDtype
 from pandas.core.arrays.string_ import (
     StringArray,
     StringDtype,
@@ -312,6 +317,20 @@ def _new_Index(cls, d):
     return cls.__new__(cls, **d)
 
 
+slice_type = Optional[
+    Union[
+        str,
+        IntegerDtype,
+        FloatingDtype,
+        DatetimeTZDtype,
+        CategoricalDtype,
+        PeriodDtype,
+        IntervalDtype,
+        abc.Hashable,
+    ]
+]
+
+
 class Index(IndexOpsMixin, PandasObject):
     """
     Immutable sequence used for indexing and alignment.
@@ -1087,7 +1106,7 @@ def view(self, cls=None):
             result._id = self._id
         return result
 
-    def astype(self, dtype, copy: bool = True):
+    def astype(self, dtype: Dtype, copy: bool = True):
         """
         Create an Index with values cast to dtypes.
 
@@ -2957,7 +2976,7 @@ def _dti_setop_align_tzs(self, other: Index, setop: str_t) -> tuple[Index, Index
         return self, other
 
     @final
-    def union(self, other, sort=None):
+    def union(self, other, sort: bool | None = None):
         """
         Form the union of two Index objects.
 
@@ -3334,7 +3353,7 @@ def _intersection_via_get_indexer(
         return result
 
     @final
-    def difference(self, other, sort=None):
+    def difference(self, other, sort: bool | None = None):
         """
         Return a new Index with elements of index not in `other`.
 
@@ -3420,7 +3439,12 @@ def _wrap_difference_result(self, other, result):
         # We will override for MultiIndex to handle empty results
         return self._wrap_setop_result(other, result)
 
-    def symmetric_difference(self, other, result_name=None, sort=None):
+    def symmetric_difference(
+        self,
+        other,
+        result_name: abc.Hashable | None = None,
+        sort: bool | None = None,
+    ):
         """
         Compute the symmetric difference of two Index objects.
 
@@ -6389,7 +6413,7 @@ def _transform_index(self, func, *, level=None) -> Index:
             items = [func(x) for x in self]
             return Index(items, name=self.name, tupleize_cols=False)
 
-    def isin(self, values, level=None) -> npt.NDArray[np.bool_]:
+    def isin(self, values, level: np.str_ | int | None = None) -> npt.NDArray[np.bool_]:
         """
         Return a boolean array where the index values are in `values`.
 
@@ -6687,7 +6711,12 @@ def get_slice_bound(self, label, side: Literal["left", "right"]) -> int:
             else:
                 return slc
 
-    def slice_locs(self, start=None, end=None, step=None) -> tuple[int, int]:
+    def slice_locs(
+        self,
+        start: slice_type = None,
+        end: slice_type = None,
+        step: int | None = None,
+    ) -> tuple[int, int]:
         """
         Compute slice locations for input labels.
 
@@ -6781,7 +6810,7 @@ def slice_locs(self, start=None, end=None, step=None) -> tuple[int, int]:
 
         return start_slice, end_slice
 
-    def delete(self, loc) -> Self:
+    def delete(self, loc: int | list[int] | np.ndarray) -> Self:
         """
         Make new Index with passed location(-s) deleted.
 
@@ -7227,7 +7256,9 @@ def _maybe_disable_logical_methods(self, opname: str_t) -> None:
             raise TypeError(f"cannot perform {opname} with {type(self).__name__}")
 
     @Appender(IndexOpsMixin.argmin.__doc__)
-    def argmin(self, axis=None, skipna: bool = True, *args, **kwargs) -> int:
+    def argmin(
+        self, axis: int | None = None, skipna: bool = True, *args, **kwargs
+    ) -> int:
         nv.validate_argmin(args, kwargs)
         nv.validate_minmax_axis(axis)
 
@@ -7240,7 +7271,9 @@ def argmin(self, axis=None, skipna: bool = True, *args, **kwargs) -> int:
         return super().argmin(skipna=skipna)
 
     @Appender(IndexOpsMixin.argmax.__doc__)
-    def argmax(self, axis=None, skipna: bool = True, *args, **kwargs) -> int:
+    def argmax(
+        self, axis: int | None = None, skipna: bool = True, *args, **kwargs
+    ) -> int:
         nv.validate_argmax(args, kwargs)
         nv.validate_minmax_axis(axis)
 
@@ -7251,7 +7284,7 @@ def argmax(self, axis=None, skipna: bool = True, *args, **kwargs) -> int:
                 raise ValueError("Encountered all NA values")
         return super().argmax(skipna=skipna)
 
-    def min(self, axis=None, skipna: bool = True, *args, **kwargs):
+    def min(self, axis: int | None = None, skipna: bool = True, *args, **kwargs):
         """
         Return the minimum value of the Index.
 
@@ -7314,7 +7347,7 @@ def min(self, axis=None, skipna: bool = True, *args, **kwargs):
 
         return nanops.nanmin(self._values, skipna=skipna)
 
-    def max(self, axis=None, skipna: bool = True, *args, **kwargs):
+    def max(self, axis: int | None = None, skipna: bool = True, *args, **kwargs):
         """
         Return the maximum value of the Index.