Skip to content

TYP: indexes #43273

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 29, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 12 additions & 8 deletions pandas/core/indexes/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -2498,7 +2498,7 @@ def __reduce__(self):
"""The expected NA value to use with this index."""

@cache_readonly
def _isnan(self) -> np.ndarray:
def _isnan(self) -> npt.NDArray[np.bool_]:
"""
Return if each value is NaN.
"""
Expand All @@ -2521,7 +2521,7 @@ def hasnans(self) -> bool:
return False

@final
def isna(self) -> np.ndarray:
def isna(self) -> npt.NDArray[np.bool_]:
"""
Detect missing values.

Expand Down Expand Up @@ -2579,7 +2579,7 @@ def isna(self) -> np.ndarray:
isnull = isna

@final
def notna(self) -> np.ndarray:
def notna(self) -> npt.NDArray[np.bool_]:
"""
Detect existing (non-missing) values.

Expand Down Expand Up @@ -2764,7 +2764,9 @@ def drop_duplicates(self: _IndexT, keep: str_t | bool = "first") -> _IndexT:

return super().drop_duplicates(keep=keep)

def duplicated(self, keep: Literal["first", "last", False] = "first") -> np.ndarray:
def duplicated(
self, keep: Literal["first", "last", False] = "first"
) -> npt.NDArray[np.bool_]:
"""
Indicate duplicate index values.

Expand Down Expand Up @@ -5495,7 +5497,7 @@ def _get_indexer_strict(self, key, axis_name: str_t) -> tuple[Index, np.ndarray]

return keyarr, indexer

def _raise_if_missing(self, key, indexer, axis_name: str_t):
def _raise_if_missing(self, key, indexer, axis_name: str_t) -> None:
"""
Check that indexer can be used to return a result.

Expand Down Expand Up @@ -6119,7 +6121,9 @@ def get_slice_bound(self, label, side: str_t, kind=no_default) -> int:
else:
return slc

def slice_locs(self, start=None, end=None, step=None, kind=no_default):
def slice_locs(
self, start=None, end=None, step=None, kind=no_default
) -> tuple[int, int]:
"""
Compute slice locations for input labels.

Expand Down Expand Up @@ -6491,7 +6495,7 @@ def all(self, *args, **kwargs):
return np.all(self.values) # type: ignore[arg-type]

@final
def _maybe_disable_logical_methods(self, opname: str_t):
def _maybe_disable_logical_methods(self, opname: str_t) -> None:
"""
raise if this Index subclass does not support any or all.
"""
Expand Down Expand Up @@ -6606,7 +6610,7 @@ def _deprecated_arg(self, value, name: str_t, methodname: str_t) -> None:
)


def ensure_index_from_sequences(sequences, names=None):
def ensure_index_from_sequences(sequences, names=None) -> Index:
"""
Construct an index from sequences of data.

Expand Down
7 changes: 5 additions & 2 deletions pandas/core/indexes/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,10 @@

import numpy as np

from pandas._typing import ArrayLike
from pandas._typing import (
ArrayLike,
npt,
)
from pandas.compat.numpy import function as nv
from pandas.util._decorators import (
cache_readonly,
Expand Down Expand Up @@ -437,7 +440,7 @@ def astype(self, dtype, copy: bool = True) -> Index:
return Index(new_values, dtype=new_values.dtype, name=self.name, copy=False)

@cache_readonly
def _isnan(self) -> np.ndarray:
def _isnan(self) -> npt.NDArray[np.bool_]:
# error: Incompatible return value type (got "ExtensionArray", expected
# "ndarray")
return self._data.isna() # type: ignore[return-value]
Expand Down
33 changes: 19 additions & 14 deletions pandas/core/indexes/multi.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
DtypeObj,
Scalar,
Shape,
npt,
)
from pandas.compat.numpy import function as nv
from pandas.errors import (
Expand Down Expand Up @@ -1588,7 +1589,7 @@ def _inferred_type_levels(self) -> list[str]:
return [i.inferred_type for i in self.levels]

@doc(Index.duplicated)
def duplicated(self, keep="first") -> np.ndarray:
def duplicated(self, keep="first") -> npt.NDArray[np.bool_]:
shape = tuple(len(lev) for lev in self.levels)
ids = get_group_index(self.codes, shape, sort=False, xnull=False)

Expand Down Expand Up @@ -1842,7 +1843,7 @@ def _is_lexsorted(self) -> bool:
return self._lexsort_depth == self.nlevels

@property
def lexsort_depth(self):
def lexsort_depth(self) -> int:
warnings.warn(
"MultiIndex.is_lexsorted is deprecated as a public function, "
"users should use MultiIndex.is_monotonic_increasing instead.",
Expand Down Expand Up @@ -2152,7 +2153,7 @@ def append(self, other):
except (TypeError, IndexError):
return Index._with_infer(new_tuples)

def argsort(self, *args, **kwargs) -> np.ndarray:
def argsort(self, *args, **kwargs) -> npt.NDArray[np.intp]:
return self._values.argsort(*args, **kwargs)

@Appender(_index_shared_docs["repeat"] % _index_doc_kwargs)
Expand Down Expand Up @@ -2371,7 +2372,7 @@ def cats(level_codes):

def sortlevel(
self, level=0, ascending: bool = True, sort_remaining: bool = True
) -> tuple[MultiIndex, np.ndarray]:
) -> tuple[MultiIndex, npt.NDArray[np.intp]]:
"""
Sort MultiIndex at the requested level.

Expand All @@ -2392,7 +2393,7 @@ def sortlevel(
-------
sorted_index : pd.MultiIndex
Resulting index.
indexer : np.ndarray
indexer : np.ndarray[np.intp]
Indices of output values in original index.

Examples
Expand Down Expand Up @@ -2493,7 +2494,7 @@ def _wrap_reindex_result(self, target, indexer, preserve_names: bool):
target = self._maybe_preserve_names(target, preserve_names)
return target

def _maybe_preserve_names(self, target: Index, preserve_names: bool):
def _maybe_preserve_names(self, target: Index, preserve_names: bool) -> Index:
if (
preserve_names
and target.nlevels == self.nlevels
Expand All @@ -2506,7 +2507,7 @@ def _maybe_preserve_names(self, target: Index, preserve_names: bool):
# --------------------------------------------------------------------
# Indexing Methods

def _check_indexing_error(self, key):
def _check_indexing_error(self, key) -> None:
if not is_hashable(key) or is_iterator(key):
# We allow tuples if they are hashable, whereas other Index
# subclasses require scalar.
Expand Down Expand Up @@ -2541,7 +2542,9 @@ def _get_values_for_loc(self, series: Series, loc, key):
new_ser = series._constructor(new_values, index=new_index, name=series.name)
return new_ser.__finalize__(series)

def _get_indexer_strict(self, key, axis_name: str) -> tuple[Index, np.ndarray]:
def _get_indexer_strict(
self, key, axis_name: str
) -> tuple[Index, npt.NDArray[np.intp]]:

keyarr = key
if not isinstance(keyarr, Index):
Expand All @@ -2555,7 +2558,7 @@ def _get_indexer_strict(self, key, axis_name: str) -> tuple[Index, np.ndarray]:

return super()._get_indexer_strict(key, axis_name)

def _raise_if_missing(self, key, indexer, axis_name: str):
def _raise_if_missing(self, key, indexer, axis_name: str) -> None:
keyarr = key
if not isinstance(key, Index):
keyarr = com.asarray_tuplesafe(key)
Expand All @@ -2575,7 +2578,7 @@ def _raise_if_missing(self, key, indexer, axis_name: str):
else:
return super()._raise_if_missing(key, indexer, axis_name)

def _get_indexer_level_0(self, target) -> np.ndarray:
def _get_indexer_level_0(self, target) -> npt.NDArray[np.intp]:
"""
Optimized equivalent to `self.get_level_values(0).get_indexer_for(target)`.
"""
Expand Down Expand Up @@ -2640,7 +2643,9 @@ def get_slice_bound(
label = (label,)
return self._partial_tup_index(label, side=side)

def slice_locs(self, start=None, end=None, step=None, kind=lib.no_default):
def slice_locs(
self, start=None, end=None, step=None, kind=lib.no_default
) -> tuple[int, int]:
"""
For an ordered MultiIndex, compute the slice locations for input
labels.
Expand Down Expand Up @@ -3614,7 +3619,7 @@ def _maybe_match_names(self, other):
names.append(None)
return names

def _wrap_intersection_result(self, other, result):
def _wrap_intersection_result(self, other, result) -> MultiIndex:
_, result_names = self._convert_can_do_setop(other)

if len(result) == 0:
Expand All @@ -3627,7 +3632,7 @@ def _wrap_intersection_result(self, other, result):
else:
return MultiIndex.from_arrays(zip(*result), sortorder=0, names=result_names)

def _wrap_difference_result(self, other, result):
def _wrap_difference_result(self, other, result) -> MultiIndex:
_, result_names = self._convert_can_do_setop(other)

if len(result) == 0:
Expand Down Expand Up @@ -3738,7 +3743,7 @@ def delete(self, loc) -> MultiIndex:
)

@doc(Index.isin)
def isin(self, values, level=None) -> np.ndarray:
def isin(self, values, level=None) -> npt.NDArray[np.bool_]:
if level is None:
values = MultiIndex.from_tuples(values, names=self.names)._values
return algos.isin(self._values, values)
Expand Down