diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 7eaadecbd6491..07862e0b9bb48 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -1,4 +1,6 @@ -from typing import Any, Optional, Sequence, Type, TypeVar +from __future__ import annotations + +from typing import Any, Optional, Sequence, Type, TypeVar, Union import numpy as np @@ -212,7 +214,9 @@ def __setitem__(self, key, value): def _validate_setitem_value(self, value): return value - def __getitem__(self, key): + def __getitem__( + self: NDArrayBackedExtensionArrayT, key: Union[int, slice, np.ndarray] + ) -> Union[NDArrayBackedExtensionArrayT, Any]: if lib.is_integer(key): # fast-path result = self._ndarray[key] diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index 0968545a6b8a4..a0a51495791d1 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -6,6 +6,8 @@ This is an experimental API and subject to breaking changes without warning. """ +from __future__ import annotations + import operator from typing import ( Any, @@ -254,8 +256,9 @@ def _from_factorized(cls, values, original): # Must be a Sequence # ------------------------------------------------------------------------ - def __getitem__(self, item): - # type (Any) -> Any + def __getitem__( + self, item: Union[int, slice, np.ndarray] + ) -> Union[ExtensionArray, Any]: """ Select a subset of self. @@ -661,7 +664,7 @@ def dropna(self): """ return self[~self.isna()] - def shift(self, periods: int = 1, fill_value: object = None) -> "ExtensionArray": + def shift(self, periods: int = 1, fill_value: object = None) -> ExtensionArray: """ Shift values by desired number. @@ -831,7 +834,7 @@ def _values_for_factorize(self) -> Tuple[np.ndarray, Any]: """ return self.astype(object), np.nan - def factorize(self, na_sentinel: int = -1) -> Tuple[np.ndarray, "ExtensionArray"]: + def factorize(self, na_sentinel: int = -1) -> Tuple[np.ndarray, ExtensionArray]: """ Encode the extension array as an enumerated type. @@ -940,7 +943,7 @@ def take( *, allow_fill: bool = False, fill_value: Any = None, - ) -> "ExtensionArray": + ) -> ExtensionArray: """ Take elements from an array. @@ -1109,7 +1112,7 @@ def _formatter(self, boxed: bool = False) -> Callable[[Any], Optional[str]]: # Reshaping # ------------------------------------------------------------------------ - def transpose(self, *axes) -> "ExtensionArray": + def transpose(self, *axes) -> ExtensionArray: """ Return a transposed view on this array. @@ -1119,10 +1122,10 @@ def transpose(self, *axes) -> "ExtensionArray": return self[:] @property - def T(self) -> "ExtensionArray": + def T(self) -> ExtensionArray: return self.transpose() - def ravel(self, order="C") -> "ExtensionArray": + def ravel(self, order="C") -> ExtensionArray: """ Return a flattened view on this array. diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index 0ce32fcd822e0..30674212239bf 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from datetime import datetime, timedelta import operator from typing import ( @@ -264,7 +266,9 @@ def __array__(self, dtype=None) -> np.ndarray: return np.array(list(self), dtype=object) return self._ndarray - def __getitem__(self, key): + def __getitem__( + self, key: Union[int, slice, np.ndarray] + ) -> Union[DatetimeLikeArrayMixin, DTScalarOrNaT]: """ This getitem defers to the underlying array, which by-definition can only handle list-likes, slices, and integer scalars diff --git a/pandas/core/arrays/datetimes.py b/pandas/core/arrays/datetimes.py index a05dc717f83c1..057162fedaa98 100644 --- a/pandas/core/arrays/datetimes.py +++ b/pandas/core/arrays/datetimes.py @@ -1,5 +1,5 @@ from datetime import datetime, time, timedelta, tzinfo -from typing import Optional, Union +from typing import Optional, Union, cast import warnings import numpy as np @@ -444,9 +444,11 @@ def _generate_range( ) if not left_closed and len(index) and index[0] == start: - index = index[1:] + # TODO: overload DatetimeLikeArrayMixin.__getitem__ + index = cast(DatetimeArray, index[1:]) if not right_closed and len(index) and index[-1] == end: - index = index[:-1] + # TODO: overload DatetimeLikeArrayMixin.__getitem__ + index = cast(DatetimeArray, index[:-1]) dtype = tz_to_dtype(tz) return cls._simple_new(index.asi8, freq=freq, dtype=dtype) diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py index a4b88427ceb05..caed932cd7857 100644 --- a/pandas/core/arrays/masked.py +++ b/pandas/core/arrays/masked.py @@ -1,4 +1,6 @@ -from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Type, TypeVar +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Optional, Sequence, Tuple, Type, TypeVar, Union import numpy as np @@ -56,7 +58,7 @@ def itemsize(self) -> int: return self.numpy_dtype.itemsize @classmethod - def construct_array_type(cls) -> Type["BaseMaskedArray"]: + def construct_array_type(cls) -> Type[BaseMaskedArray]: """ Return the array type associated with this dtype. @@ -100,7 +102,9 @@ def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False): def dtype(self) -> BaseMaskedDtype: raise AbstractMethodError(self) - def __getitem__(self, item): + def __getitem__( + self, item: Union[int, slice, np.ndarray] + ) -> Union[BaseMaskedArray, Any]: if is_integer(item): if self._mask[item]: return self.dtype.na_value diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py index ec96a0d502d3f..ad8f212aa20ea 100644 --- a/pandas/core/groupby/groupby.py +++ b/pandas/core/groupby/groupby.py @@ -1671,10 +1671,10 @@ def first(self, numeric_only: bool = False, min_count: int = -1): def first_compat(obj: FrameOrSeries, axis: int = 0): def first(x: Series): """Helper function for first item that isn't NA.""" - x = x.array[notna(x.array)] - if len(x) == 0: + arr = x.array[notna(x.array)] + if not len(arr): return np.nan - return x[0] + return arr[0] if isinstance(obj, DataFrame): return obj.apply(first, axis=axis) @@ -1695,10 +1695,10 @@ def last(self, numeric_only: bool = False, min_count: int = -1): def last_compat(obj: FrameOrSeries, axis: int = 0): def last(x: Series): """Helper function for last item that isn't NA.""" - x = x.array[notna(x.array)] - if len(x) == 0: + arr = x.array[notna(x.array)] + if not len(arr): return np.nan - return x[-1] + return arr[-1] if isinstance(obj, DataFrame): return obj.apply(last, axis=axis) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index cb5641a74e60b..1160f374643c7 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -3212,8 +3212,14 @@ def _get_nearest_indexer(self, target: "Index", limit, tolerance) -> np.ndarray: right_indexer = self.get_indexer(target, "backfill", limit=limit) target_values = target._values - left_distances = np.abs(self._values[left_indexer] - target_values) - right_distances = np.abs(self._values[right_indexer] - target_values) + # error: Unsupported left operand type for - ("ExtensionArray") + left_distances = np.abs( + self._values[left_indexer] - target_values # type: ignore[operator] + ) + # error: Unsupported left operand type for - ("ExtensionArray") + right_distances = np.abs( + self._values[right_indexer] - target_values # type: ignore[operator] + ) op = operator.lt if self.is_monotonic_increasing else operator.le indexer = np.where( @@ -3232,7 +3238,8 @@ def _filter_indexer_tolerance( indexer: np.ndarray, tolerance, ) -> np.ndarray: - distance = abs(self._values[indexer] - target) + # error: Unsupported left operand type for - ("ExtensionArray") + distance = abs(self._values[indexer] - target) # type: ignore[operator] indexer = np.where(distance <= tolerance, indexer, -1) return indexer @@ -3436,6 +3443,7 @@ def reindex(self, target, method=None, level=None, limit=None, tolerance=None): target = ensure_has_len(target) # target may be an iterator if not isinstance(target, Index) and len(target) == 0: + values: Union[range, ExtensionArray, np.ndarray] if isinstance(self, ABCRangeIndex): values = range(0) else: @@ -4528,8 +4536,9 @@ def asof_locs(self, where: "Index", mask) -> np.ndarray: result = np.arange(len(self))[mask].take(locs) - first = mask.argmax() - result[(locs == 0) & (where._values < self._values[first])] = -1 + # TODO: overload return type of ExtensionArray.__getitem__ + first_value = cast(Any, self._values[mask.argmax()]) + result[(locs == 0) & (where._values < first_value)] = -1 return result diff --git a/pandas/core/indexes/period.py b/pandas/core/indexes/period.py index 44c20ad0de848..aaccef61fbaba 100644 --- a/pandas/core/indexes/period.py +++ b/pandas/core/indexes/period.py @@ -1,5 +1,5 @@ from datetime import datetime, timedelta -from typing import Any +from typing import Any, cast import numpy as np @@ -694,7 +694,10 @@ def difference(self, other, sort=None): if self.equals(other): # pass an empty PeriodArray with the appropriate dtype - return type(self)._simple_new(self._data[:0], name=self.name) + + # TODO: overload DatetimeLikeArrayMixin.__getitem__ + values = cast(PeriodArray, self._data[:0]) + return type(self)._simple_new(values, name=self.name) if is_object_dtype(other): return self.astype(object).difference(other).astype(self.dtype)