TYP: indexes

pandas/core/arrays/datetimes.py

@@ -192,7 +192,7 @@ class DatetimeArray(dtl.TimelikeOps, dtl.DatelikeOps):
     _infer_matches = ("datetime", "datetime64", "date")
 
     # define my properties & methods for delegation
-    _bool_ops = [
+    _bool_ops: list[str] = [
         "is_month_start",
         "is_month_end",
         "is_quarter_start",
@@ -201,8 +201,8 @@ class DatetimeArray(dtl.TimelikeOps, dtl.DatelikeOps):
         "is_year_end",
         "is_leap_year",
     ]
-    _object_ops = ["freq", "tz"]
-    _field_ops = [
+    _object_ops: list[str] = ["freq", "tz"]
+    _field_ops: list[str] = [
         "year",
         "month",
         "day",
@@ -222,9 +222,9 @@ class DatetimeArray(dtl.TimelikeOps, dtl.DatelikeOps):
         "microsecond",
         "nanosecond",
     ]
-    _other_ops = ["date", "time", "timetz"]
-    _datetimelike_ops = _field_ops + _object_ops + _bool_ops + _other_ops
-    _datetimelike_methods = [
+    _other_ops: list[str] = ["date", "time", "timetz"]
+    _datetimelike_ops: list[str] = _field_ops + _object_ops + _bool_ops + _other_ops
+    _datetimelike_methods: list[str] = [
         "to_period",
         "tz_localize",
         "tz_convert",

pandas/core/arrays/interval.py

@@ -89,7 +89,7 @@
 
 IntervalArrayT = TypeVar("IntervalArrayT", bound="IntervalArray")
 
-_interval_shared_docs = {}
+_interval_shared_docs: dict[str, str] = {}
 
 _shared_docs_kwargs = {
     "klass": "IntervalArray",

pandas/core/arrays/period.py

@@ -6,7 +6,6 @@
     TYPE_CHECKING,
     Any,
     Callable,
-    List,
     Optional,
     Sequence,
     Type,
@@ -160,10 +159,10 @@ class PeriodArray(PeriodMixin, dtl.DatelikeOps):
     _infer_matches = ("period",)
 
     # Names others delegate to us
-    _other_ops: List[str] = []
-    _bool_ops = ["is_leap_year"]
-    _object_ops = ["start_time", "end_time", "freq"]
-    _field_ops = [
+    _other_ops: list[str] = []
+    _bool_ops: list[str] = ["is_leap_year"]
+    _object_ops: list[str] = ["start_time", "end_time", "freq"]
+    _field_ops: list[str] = [
         "year",
         "month",
         "day",

pandas/core/generic.py

@@ -1822,9 +1822,9 @@ def _drop_labels_or_levels(self, keys, axis: int = 0):
 
         Parameters
         ----------
-        keys: str or list of str
+        keys : str or list of str
             labels or levels to drop
-        axis: int, default 0
+        axis : int, default 0
             Axis that levels are associated with (0 for index, 1 for columns)
 
         Returns

pandas/core/groupby/base.py

@@ -3,6 +3,8 @@
 hold the allowlist of methods that are exposed on the
 SeriesGroupBy and the DataFrameGroupBy objects.
 """
+from __future__ import annotations
+
 import collections
 from typing import List
 
@@ -106,12 +108,16 @@ def _gotitem(self, key, ndim, subset=None):
     | plotting_methods
 )
 
-series_apply_allowlist = (
+series_apply_allowlist: frozenset[str] = (
     common_apply_allowlist
-    | {"nlargest", "nsmallest", "is_monotonic_increasing", "is_monotonic_decreasing"}
+    | frozenset(
+        {"nlargest", "nsmallest", "is_monotonic_increasing", "is_monotonic_decreasing"}
+    )
 ) | frozenset(["dtype", "unique"])
 
-dataframe_apply_allowlist = common_apply_allowlist | frozenset(["dtypes", "corrwith"])
+dataframe_apply_allowlist: frozenset[str] = common_apply_allowlist | frozenset(
+    ["dtypes", "corrwith"]
+)
 
 # cythonized transformations or canned "agg+broadcast", which do not
 # require postprocessing of the result by transform.

pandas/core/indexes/base.py

@@ -179,15 +179,15 @@
 
 _unsortable_types = frozenset(("mixed", "mixed-integer"))
 
-_index_doc_kwargs = {
+_index_doc_kwargs: dict[str, str] = {
     "klass": "Index",
     "inplace": "",
     "target_klass": "Index",
     "raises_section": "",
     "unique": "Index",
     "duplicated": "np.ndarray",
 }
-_index_shared_docs = {}
+_index_shared_docs: dict[str, str] = {}
 str_t = str
 
 
@@ -1167,7 +1167,7 @@ def _format_with_header(
         return header + result
 
     @final
-    def to_native_types(self, slicer=None, **kwargs):
+    def to_native_types(self, slicer=None, **kwargs) -> np.ndarray:
         """
         Format specified values of `self` and return them.
 
@@ -4410,7 +4410,7 @@ def memory_usage(self, deep: bool = False) -> int:
         return result
 
     @final
-    def where(self, cond, other=None):
+    def where(self, cond, other=None) -> Index:
         """
         Replace values where the condition is False.
 
@@ -4626,7 +4626,7 @@ def _can_hold_identifiers_and_holds_name(self, name) -> bool:
             return name in self
         return False
 
-    def append(self, other):
+    def append(self, other) -> Index:
         """
         Append a collection of Index options together.
 
@@ -4636,7 +4636,7 @@ def append(self, other):
 
         Returns
         -------
-        appended : Index
+        Index
         """
         to_concat = [self]
 
@@ -4866,7 +4866,7 @@ def asof(self, label):
                 loc = loc.indices(len(self))[-1]
             return self[loc]
 
-    def asof_locs(self, where: Index, mask) -> np.ndarray:
+    def asof_locs(self, where: Index, mask: np.ndarray) -> np.ndarray:
         """
         Return the locations (indices) of labels in the index.
 
@@ -4883,24 +4883,26 @@ def asof_locs(self, where: Index, mask) -> np.ndarray:
         ----------
         where : Index
             An Index consisting of an array of timestamps.
-        mask : array-like
+        mask : np.ndarray[bool]
             Array of booleans denoting where values in the original
             data are not NA.
 
         Returns
         -------
-        numpy.ndarray
+        np.ndarray[np.intp]
             An array of locations (indices) of the labels from the Index
             which correspond to the return values of the `asof` function
             for every element in `where`.
         """
         locs = self._values[mask].searchsorted(where._values, side="right")
         locs = np.where(locs > 0, locs - 1, 0)
 
-        result = np.arange(len(self))[mask].take(locs)
+        result = np.arange(len(self), dtype=np.intp)[mask].take(locs)
 
         # TODO: overload return type of ExtensionArray.__getitem__
-        first_value = cast(Any, self._values[mask.argmax()])
+        # error: Invalid index type "signedinteger[Any]" for
+        # "Union[ExtensionArray, ndarray]"; expected type "Union[int, slice, ndarray]"
+        first_value = cast(Any, self._values[mask.argmax()])  # type: ignore[index]
         result[(locs == 0) & (where._values < first_value)] = -1
 
         return result
@@ -5070,7 +5072,7 @@ def argsort(self, *args, **kwargs) -> np.ndarray:
 
         Returns
         -------
-        numpy.ndarray
+        np.ndarray[np.intp]
             Integer indices that would sort the index if used as
             an indexer.
 
@@ -5858,7 +5860,7 @@ def delete(self, loc) -> Index:
         Returns
         -------
         Index
-            New Index with passed location(-s) deleted.
+            Will be same type as self, except for RangeIndex.
 
         See Also
         --------
@@ -6374,8 +6376,8 @@ def _maybe_cast_data_without_dtype(subarr):
 
     elif inferred == "interval":
         try:
-            data = IntervalArray._from_sequence(subarr, copy=False)
-            return data
+            ia_data = IntervalArray._from_sequence(subarr, copy=False)
+            return ia_data
         except (ValueError, TypeError):
             # GH27172: mixed closed Intervals --> object dtype
             pass

pandas/core/indexes/category.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import (
     Any,
     Hashable,
@@ -50,7 +52,7 @@
     inherit_names,
 )
 
-_index_doc_kwargs = dict(ibase._index_doc_kwargs)
+_index_doc_kwargs: dict[str, str] = dict(ibase._index_doc_kwargs)
 _index_doc_kwargs.update({"target_klass": "CategoricalIndex"})
 
 
@@ -217,9 +219,9 @@ def __new__(
         categories=None,
         ordered=None,
         dtype: Optional[Dtype] = None,
-        copy=False,
-        name=None,
-    ):
+        copy: bool = False,
+        name: Hashable = None,
+    ) -> CategoricalIndex:
 
         name = maybe_extract_name(name, data, cls)
 
@@ -239,7 +241,7 @@ def _shallow_copy(
         self,
         values: Categorical,
         name: Hashable = no_default,
-    ):
+    ) -> CategoricalIndex:
         name = self._name if name is no_default else name
 
         if values is not None:

pandas/core/indexes/datetimes.py

@@ -10,6 +10,7 @@
 import operator
 from typing import (
     TYPE_CHECKING,
+    Hashable,
     Optional,
     Tuple,
 )
@@ -318,15 +319,15 @@ def __new__(
         data=None,
         freq=lib.no_default,
         tz=None,
-        normalize=False,
+        normalize: bool = False,
         closed=None,
         ambiguous="raise",
-        dayfirst=False,
-        yearfirst=False,
+        dayfirst: bool = False,
+        yearfirst: bool = False,
         dtype: Optional[Dtype] = None,
-        copy=False,
-        name=None,
-    ):
+        copy: bool = False,
+        name: Hashable = None,
+    ) -> DatetimeIndex:
 
         if is_scalar(data):
             raise TypeError(
@@ -641,7 +642,7 @@ def _validate_partial_date_slice(self, reso: Resolution):
             # _parsed_string_to_bounds allows it.
             raise KeyError
 
-    def _deprecate_mismatched_indexing(self, key):
+    def _deprecate_mismatched_indexing(self, key) -> None:
         # GH#36148
         # we get here with isinstance(key, self._data._recognized_scalars)
         try:
@@ -855,7 +856,7 @@ def inferred_type(self) -> str:
         # sure we can't have ambiguous indexing
         return "datetime64"
 
-    def indexer_at_time(self, time, asof=False):
+    def indexer_at_time(self, time, asof: bool = False) -> np.ndarray:
         """
         Return index locations of values at particular time of day
         (e.g. 9:30AM).
@@ -869,7 +870,7 @@ def indexer_at_time(self, time, asof=False):
 
         Returns
         -------
-        values_at_time : array of integers
+        np.ndarray[np.intp]
 
         See Also
         --------
@@ -895,8 +896,8 @@ def indexer_at_time(self, time, asof=False):
         return (micros == time_micros).nonzero()[0]
 
     def indexer_between_time(
-        self, start_time, end_time, include_start=True, include_end=True
-    ):
+        self, start_time, end_time, include_start: bool = True, include_end: bool = True
+    ) -> np.ndarray:
         """
         Return index locations of values between particular times of day
         (e.g., 9:00-9:30AM).
@@ -912,7 +913,7 @@ def indexer_between_time(
 
         Returns
         -------
-        values_between_time : array of integers
+        np.ndarray[np.intp]
 
         See Also
         --------
@@ -952,8 +953,8 @@ def date_range(
     periods=None,
     freq=None,
     tz=None,
-    normalize=False,
-    name=None,
+    normalize: bool = False,
+    name: Hashable = None,
     closed=None,
     **kwargs,
 ) -> DatetimeIndex:
@@ -1124,8 +1125,8 @@ def bdate_range(
     periods: Optional[int] = None,
     freq="B",
     tz=None,
-    normalize=True,
-    name=None,
+    normalize: bool = True,
+    name: Hashable = None,
     weekmask=None,
     holidays=None,
     closed=None,

pandas/core/indexes/extension.py

@@ -351,7 +351,7 @@ def map(self, mapper, na_action=None):
             return self.astype(object).map(mapper)
 
     @doc(Index.astype)
-    def astype(self, dtype, copy=True):
+    def astype(self, dtype, copy: bool = True) -> Index:
         dtype = pandas_dtype(dtype)
         if is_dtype_equal(self.dtype, dtype):
             if not copy:
@@ -410,7 +410,7 @@ def _simple_new(
     def _get_engine_target(self) -> np.ndarray:
         return self._data._ndarray
 
-    def insert(self: _T, loc: int, item) -> _T:
+    def insert(self: _T, loc: int, item) -> Index:
         """
         Make new Index inserting new item at location. Follows
         Python list.append semantics for negative values.