diff --git a/pandas/_typing.py b/pandas/_typing.py index 93d49497a85e0..5077e659410e3 100644 --- a/pandas/_typing.py +++ b/pandas/_typing.py @@ -206,10 +206,16 @@ # indexing # PositionalIndexer -> valid 1D positional indexer, e.g. can pass # to ndarray.__getitem__ +# ScalarIndexer is for a single value as the index +# SequenceIndexer is for list like or slices (but not tuples) +# PositionalIndexerTuple is extends the PositionalIndexer for 2D arrays +# These are used in various __getitem__ overloads # TODO: add Ellipsis, see # https://github.com/python/typing/issues/684#issuecomment-548203158 # https://bugs.python.org/issue41810 -PositionalIndexer = Union[int, np.integer, slice, Sequence[int], np.ndarray] -PositionalIndexer2D = Union[ - PositionalIndexer, Tuple[PositionalIndexer, PositionalIndexer] -] +# Using List[int] here rather than Sequence[int] to disallow tuples. +ScalarIndexer = Union[int, np.integer] +SequenceIndexer = Union[slice, List[int], np.ndarray] +PositionalIndexer = Union[ScalarIndexer, SequenceIndexer] +PositionalIndexerTuple = Tuple[PositionalIndexer, PositionalIndexer] +PositionalIndexer2D = Union[PositionalIndexer, PositionalIndexerTuple] diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index f13f1a418c2e9..4c7ccc2f16477 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -4,9 +4,11 @@ from typing import ( TYPE_CHECKING, Any, + Literal, Sequence, TypeVar, cast, + overload, ) import numpy as np @@ -16,6 +18,9 @@ from pandas._typing import ( F, PositionalIndexer2D, + PositionalIndexerTuple, + ScalarIndexer, + SequenceIndexer, Shape, npt, type_t, @@ -48,7 +53,6 @@ ) if TYPE_CHECKING: - from typing import Literal from pandas._typing import ( NumpySorter, @@ -205,6 +209,17 @@ def __setitem__(self, key, value): def _validate_setitem_value(self, value): return value + @overload + def __getitem__(self, key: ScalarIndexer) -> Any: + ... + + @overload + def __getitem__( + self: NDArrayBackedExtensionArrayT, + key: SequenceIndexer | PositionalIndexerTuple, + ) -> NDArrayBackedExtensionArrayT: + ... + def __getitem__( self: NDArrayBackedExtensionArrayT, key: PositionalIndexer2D, diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index b0b7b81d059e6..40837ccad6ac8 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -30,6 +30,8 @@ Dtype, FillnaOptions, PositionalIndexer, + ScalarIndexer, + SequenceIndexer, Shape, npt, ) @@ -298,8 +300,17 @@ def _from_factorized(cls, values, original): # ------------------------------------------------------------------------ # Must be a Sequence # ------------------------------------------------------------------------ + @overload + def __getitem__(self, item: ScalarIndexer) -> Any: + ... + + @overload + def __getitem__(self: ExtensionArrayT, item: SequenceIndexer) -> ExtensionArrayT: + ... - def __getitem__(self, item: PositionalIndexer) -> ExtensionArray | Any: + def __getitem__( + self: ExtensionArrayT, item: PositionalIndexer + ) -> ExtensionArrayT | Any: """ Select a subset of self. @@ -313,6 +324,8 @@ def __getitem__(self, item: PositionalIndexer) -> ExtensionArray | Any: * ndarray: A 1-d boolean NumPy ndarray the same length as 'self' + * list[int]: A list of int + Returns ------- item : scalar or ExtensionArray @@ -761,7 +774,7 @@ def fillna( new_values = self.copy() return new_values - def dropna(self): + def dropna(self: ExtensionArrayT) -> ExtensionArrayT: """ Return ExtensionArray without NA values. diff --git a/pandas/core/arrays/categorical.py b/pandas/core/arrays/categorical.py index b8ceef3d52e41..543b018c07ea5 100644 --- a/pandas/core/arrays/categorical.py +++ b/pandas/core/arrays/categorical.py @@ -6,6 +6,7 @@ from shutil import get_terminal_size from typing import ( TYPE_CHECKING, + Any, Hashable, Sequence, TypeVar, @@ -37,7 +38,11 @@ Dtype, NpDtype, Ordered, + PositionalIndexer2D, + PositionalIndexerTuple, Scalar, + ScalarIndexer, + SequenceIndexer, Shape, npt, type_t, @@ -2017,7 +2022,18 @@ def __repr__(self) -> str: # ------------------------------------------------------------------ - def __getitem__(self, key): + @overload + def __getitem__(self, key: ScalarIndexer) -> Any: + ... + + @overload + def __getitem__( + self: CategoricalT, + key: SequenceIndexer | PositionalIndexerTuple, + ) -> CategoricalT: + ... + + def __getitem__(self: CategoricalT, key: PositionalIndexer2D) -> CategoricalT | Any: """ Return an item. """ diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index ad3120c9c27d3..63ba9fdd59fc6 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -49,6 +49,9 @@ DtypeObj, NpDtype, PositionalIndexer2D, + PositionalIndexerTuple, + ScalarIndexer, + SequenceIndexer, npt, ) from pandas.compat.numpy import function as nv @@ -313,17 +316,33 @@ def __array__(self, dtype: NpDtype | None = None) -> np.ndarray: return np.array(list(self), dtype=object) return self._ndarray + @overload + def __getitem__(self, item: ScalarIndexer) -> DTScalarOrNaT: + ... + + @overload def __getitem__( - self, key: PositionalIndexer2D - ) -> DatetimeLikeArrayMixin | DTScalarOrNaT: + self: DatetimeLikeArrayT, + item: SequenceIndexer | PositionalIndexerTuple, + ) -> DatetimeLikeArrayT: + ... + + def __getitem__( + self: DatetimeLikeArrayT, key: PositionalIndexer2D + ) -> DatetimeLikeArrayT | DTScalarOrNaT: """ This getitem defers to the underlying array, which by-definition can only handle list-likes, slices, and integer scalars """ - result = super().__getitem__(key) + # Use cast as we know we will get back a DatetimeLikeArray or DTScalar + result = cast( + Union[DatetimeLikeArrayT, DTScalarOrNaT], super().__getitem__(key) + ) if lib.is_scalar(result): return result - + else: + # At this point we know the result is an array. + result = cast(DatetimeLikeArrayT, result) result._freq = self._get_getitem_freq(key) return result @@ -1768,11 +1787,7 @@ def factorize(self, na_sentinel=-1, sort: bool = False): uniques = self.copy() # TODO: copy or view? if sort and self.freq.n < 0: codes = codes[::-1] - # TODO: overload __getitem__, a slice indexer returns same type as self - # error: Incompatible types in assignment (expression has type - # "Union[DatetimeLikeArrayMixin, Union[Any, Any]]", variable - # has type "TimelikeOps") - uniques = uniques[::-1] # type: ignore[assignment] + uniques = uniques[::-1] return codes, uniques # FIXME: shouldn't get here; we are ignoring sort return super().factorize(na_sentinel=na_sentinel) diff --git a/pandas/core/arrays/datetimes.py b/pandas/core/arrays/datetimes.py index e82b81d55807d..823103181bb82 100644 --- a/pandas/core/arrays/datetimes.py +++ b/pandas/core/arrays/datetimes.py @@ -9,7 +9,6 @@ from typing import ( TYPE_CHECKING, Literal, - cast, overload, ) import warnings @@ -478,11 +477,9 @@ def _generate_range( index = cls._simple_new(arr, freq=None, dtype=dtype) if not left_closed and len(index) and index[0] == start: - # TODO: overload DatetimeLikeArrayMixin.__getitem__ - index = cast(DatetimeArray, index[1:]) + index = index[1:] if not right_closed and len(index) and index[-1] == end: - # TODO: overload DatetimeLikeArrayMixin.__getitem__ - index = cast(DatetimeArray, index[:-1]) + index = index[:-1] dtype = tz_to_dtype(tz) return cls._simple_new(index._ndarray, freq=freq, dtype=dtype) diff --git a/pandas/core/arrays/interval.py b/pandas/core/arrays/interval.py index 41998218acd7d..732bdb112b8c3 100644 --- a/pandas/core/arrays/interval.py +++ b/pandas/core/arrays/interval.py @@ -9,7 +9,9 @@ from typing import ( Sequence, TypeVar, + Union, cast, + overload, ) import numpy as np @@ -31,6 +33,9 @@ ArrayLike, Dtype, NpDtype, + PositionalIndexer, + ScalarIndexer, + SequenceIndexer, ) from pandas.compat.numpy import function as nv from pandas.util._decorators import Appender @@ -89,6 +94,7 @@ ) IntervalArrayT = TypeVar("IntervalArrayT", bound="IntervalArray") +IntervalOrNA = Union[Interval, float] _interval_shared_docs: dict[str, str] = {} @@ -635,7 +641,17 @@ def __iter__(self): def __len__(self) -> int: return len(self._left) - def __getitem__(self, key): + @overload + def __getitem__(self, key: ScalarIndexer) -> IntervalOrNA: + ... + + @overload + def __getitem__(self: IntervalArrayT, key: SequenceIndexer) -> IntervalArrayT: + ... + + def __getitem__( + self: IntervalArrayT, key: PositionalIndexer + ) -> IntervalArrayT | IntervalOrNA: key = check_array_indexer(self, key) left = self._left[key] right = self._right[key] @@ -1633,10 +1649,11 @@ def _from_combined(self, combined: np.ndarray) -> IntervalArray: return self._shallow_copy(left=new_left, right=new_right) def unique(self) -> IntervalArray: - # Invalid index type "Tuple[slice, int]" for "Union[ExtensionArray, - # ndarray[Any, Any]]"; expected type "Union[int, integer[Any], slice, - # Sequence[int], ndarray[Any, Any]]" - nc = unique(self._combined.view("complex128")[:, 0]) # type: ignore[index] + # No overload variant of "__getitem__" of "ExtensionArray" matches argument + # type "Tuple[slice, int]" + nc = unique( + self._combined.view("complex128")[:, 0] # type: ignore[call-overload] + ) nc = nc[:, None] return self._from_combined(nc) diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py index cccfd58aa914d..877babe4f18e8 100644 --- a/pandas/core/arrays/masked.py +++ b/pandas/core/arrays/masked.py @@ -20,6 +20,8 @@ NpDtype, PositionalIndexer, Scalar, + ScalarIndexer, + SequenceIndexer, npt, type_t, ) @@ -139,7 +141,17 @@ def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False): def dtype(self) -> BaseMaskedDtype: raise AbstractMethodError(self) - def __getitem__(self, item: PositionalIndexer) -> BaseMaskedArray | Any: + @overload + def __getitem__(self, item: ScalarIndexer) -> Any: + ... + + @overload + def __getitem__(self: BaseMaskedArrayT, item: SequenceIndexer) -> BaseMaskedArrayT: + ... + + def __getitem__( + self: BaseMaskedArrayT, item: PositionalIndexer + ) -> BaseMaskedArrayT | Any: if is_integer(item): if self._mask[item]: return self.dtype.na_value diff --git a/pandas/core/arrays/period.py b/pandas/core/arrays/period.py index 2db1f00e237ee..84e611659b165 100644 --- a/pandas/core/arrays/period.py +++ b/pandas/core/arrays/period.py @@ -6,6 +6,7 @@ TYPE_CHECKING, Any, Callable, + Literal, Sequence, ) @@ -76,7 +77,6 @@ import pandas.core.common as com if TYPE_CHECKING: - from typing import Literal from pandas._typing import ( NumpySorter, diff --git a/pandas/core/arrays/sparse/array.py b/pandas/core/arrays/sparse/array.py index 6dce9b4475d1b..77142ef450487 100644 --- a/pandas/core/arrays/sparse/array.py +++ b/pandas/core/arrays/sparse/array.py @@ -10,8 +10,11 @@ TYPE_CHECKING, Any, Callable, + Literal, Sequence, TypeVar, + cast, + overload, ) import warnings @@ -30,7 +33,10 @@ AstypeArg, Dtype, NpDtype, + PositionalIndexer, Scalar, + ScalarIndexer, + SequenceIndexer, npt, ) from pandas.compat.numpy import function as nv @@ -81,11 +87,21 @@ import pandas.io.formats.printing as printing +# See https://github.com/python/typing/issues/684 if TYPE_CHECKING: - from typing import Literal + from enum import Enum + + class ellipsis(Enum): + Ellipsis = "..." + + Ellipsis = ellipsis.Ellipsis from pandas._typing import NumpySorter +else: + ellipsis = type(Ellipsis) + + # ---------------------------------------------------------------------------- # Array @@ -813,8 +829,21 @@ def value_counts(self, dropna: bool = True): # -------- # Indexing # -------- + @overload + def __getitem__(self, key: ScalarIndexer) -> Any: + ... + + @overload + def __getitem__( + self: SparseArrayT, + key: SequenceIndexer | tuple[int | ellipsis, ...], + ) -> SparseArrayT: + ... - def __getitem__(self, key): + def __getitem__( + self: SparseArrayT, + key: PositionalIndexer | tuple[int | ellipsis, ...], + ) -> SparseArrayT | Any: if isinstance(key, tuple): if len(key) > 1: @@ -824,6 +853,8 @@ def __getitem__(self, key): key = key[:-1] if len(key) > 1: raise IndexError("too many indices for array.") + if key[0] is Ellipsis: + raise ValueError("Cannot slice with Ellipsis") key = key[0] if is_integer(key): @@ -852,7 +883,8 @@ def __getitem__(self, key): key = check_array_indexer(self, key) if com.is_bool_indexer(key): - + # mypy doesn't know we have an array here + key = cast(np.ndarray, key) return self.take(np.arange(len(key), dtype=np.int32)[key]) elif hasattr(key, "__len__"): return self.take(key) diff --git a/pandas/core/arrays/string_arrow.py b/pandas/core/arrays/string_arrow.py index ab8599f0f05ba..4be7f4eb0c521 100644 --- a/pandas/core/arrays/string_arrow.py +++ b/pandas/core/arrays/string_arrow.py @@ -6,17 +6,24 @@ TYPE_CHECKING, Any, Sequence, + Union, cast, + overload, ) import numpy as np -from pandas._libs import lib +from pandas._libs import ( + lib, + missing as libmissing, +) from pandas._typing import ( Dtype, NpDtype, PositionalIndexer, Scalar, + ScalarIndexer, + SequenceIndexer, ) from pandas.compat import ( pa_version_under1p0, @@ -77,6 +84,8 @@ if TYPE_CHECKING: from pandas import Series +ArrowStringScalarOrNAT = Union[str, libmissing.NAType] + def _chk_pyarrow_available() -> None: if pa_version_under1p0: @@ -260,7 +269,17 @@ def _concat_same_type(cls, to_concat) -> ArrowStringArray: ) ) - def __getitem__(self, item: PositionalIndexer) -> Any: + @overload + def __getitem__(self, item: ScalarIndexer) -> ArrowStringScalarOrNAT: + ... + + @overload + def __getitem__(self: ArrowStringArray, item: SequenceIndexer) -> ArrowStringArray: + ... + + def __getitem__( + self: ArrowStringArray, item: PositionalIndexer + ) -> ArrowStringArray | ArrowStringScalarOrNAT: """Select a subset of self. Parameters diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py index 0e358e611f418..ec8775cf78571 100644 --- a/pandas/core/groupby/groupby.py +++ b/pandas/core/groupby/groupby.py @@ -2984,10 +2984,9 @@ def blk_func(values: ArrayLike) -> ArrayLike: if real_2d and values.ndim == 1: assert result.shape[1] == 1, result.shape - # error: Invalid index type "Tuple[slice, int]" for - # "Union[ExtensionArray, ndarray[Any, Any]]"; expected type - # "Union[int, integer[Any], slice, Sequence[int], ndarray[Any, Any]]" - result = result[:, 0] # type: ignore[index] + # error: No overload variant of "__getitem__" of "ExtensionArray" + # matches argument type "Tuple[slice, int]" + result = result[:, 0] # type: ignore[call-overload] if needs_mask: mask = mask[:, 0] @@ -3001,11 +3000,9 @@ def blk_func(values: ArrayLike) -> ArrayLike: if needs_2d and not real_2d: if result.ndim == 2: assert result.shape[1] == 1 - # error: Invalid index type "Tuple[slice, int]" for - # "Union[ExtensionArray, Any, ndarray[Any, Any]]"; expected - # type "Union[int, integer[Any], slice, Sequence[int], - # ndarray[Any, Any]]" - result = result[:, 0] # type: ignore[index] + # error: No overload variant of "__getitem__" of "ExtensionArray" + # matches argument type "Tuple[slice, int]" + result = result[:, 0] # type: ignore[call-overload] return result.T diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 645fab0d76a73..c73b3e99600d6 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -4803,11 +4803,7 @@ def __getitem__(self, key): result = getitem(key) if not is_scalar(result): - # error: Argument 1 to "ndim" has incompatible type "Union[ExtensionArray, - # Any]"; expected "Union[Union[int, float, complex, str, bytes, generic], - # Sequence[Union[int, float, complex, str, bytes, generic]], - # Sequence[Sequence[Any]], _SupportsArray]" - if np.ndim(result) > 1: # type: ignore[arg-type] + if np.ndim(result) > 1: deprecate_ndim_indexing(result) return result # NB: Using _constructor._simple_new would break if MultiIndex @@ -5122,13 +5118,17 @@ def asof_locs(self, where: Index, mask: np.ndarray) -> npt.NDArray[np.intp]: which correspond to the return values of the `asof` function for every element in `where`. """ - locs = self._values[mask].searchsorted(where._values, side="right") + # error: No overload variant of "searchsorted" of "ndarray" matches argument + # types "Union[ExtensionArray, ndarray[Any, Any]]", "str" + # TODO: will be fixed when ExtensionArray.searchsorted() is fixed + locs = self._values[mask].searchsorted( + where._values, side="right" # type: ignore[call-overload] + ) locs = np.where(locs > 0, locs - 1, 0) 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()]) + first_value = self._values[mask.argmax()] result[(locs == 0) & (where._values < first_value)] = -1 return result diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py index 920af5a13baba..b446dfe045e62 100644 --- a/pandas/core/indexes/extension.py +++ b/pandas/core/indexes/extension.py @@ -6,6 +6,7 @@ from typing import ( TYPE_CHECKING, Hashable, + Literal, TypeVar, overload, ) @@ -47,7 +48,6 @@ from pandas.core.ops import get_op_result_name if TYPE_CHECKING: - from typing import Literal from pandas._typing import ( NumpySorter, diff --git a/pandas/core/internals/array_manager.py b/pandas/core/internals/array_manager.py index 475bfe958ea06..f645cc81e8171 100644 --- a/pandas/core/internals/array_manager.py +++ b/pandas/core/internals/array_manager.py @@ -315,10 +315,9 @@ def apply_with_block(self: T, f, align_keys=None, swap_axis=True, **kwargs) -> T if self.ndim == 2 and arr.ndim == 2: # 2D for np.ndarray or DatetimeArray/TimedeltaArray assert len(arr) == 1 - # error: Invalid index type "Tuple[int, slice]" for - # "Union[ndarray, ExtensionArray]"; expected type - # "Union[int, slice, ndarray]" - arr = arr[0, :] # type: ignore[index] + # error: No overload variant of "__getitem__" of "ExtensionArray" + # matches argument type "Tuple[int, slice]" + arr = arr[0, :] # type: ignore[call-overload] result_arrays.append(arr) return type(self)(result_arrays, self._axes) @@ -841,10 +840,9 @@ def iset(self, loc: int | slice | np.ndarray, value: ArrayLike): assert value.shape[0] == len(self._axes[0]) for value_idx, mgr_idx in enumerate(indices): - # error: Invalid index type "Tuple[slice, int]" for - # "Union[ExtensionArray, ndarray]"; expected type - # "Union[int, slice, ndarray]" - value_arr = value[:, value_idx] # type: ignore[index] + # error: No overload variant of "__getitem__" of "ExtensionArray" matches + # argument type "Tuple[slice, int]" + value_arr = value[:, value_idx] # type: ignore[call-overload] self.arrays[mgr_idx] = value_arr return @@ -864,10 +862,9 @@ def insert(self, loc: int, item: Hashable, value: ArrayLike) -> None: value = extract_array(value, extract_numpy=True) if value.ndim == 2: if value.shape[0] == 1: - # error: Invalid index type "Tuple[int, slice]" for - # "Union[Any, ExtensionArray, ndarray]"; expected type - # "Union[int, slice, ndarray]" - value = value[0, :] # type: ignore[index] + # error: No overload variant of "__getitem__" of "ExtensionArray" + # matches argument type "Tuple[int, slice]" + value = value[0, :] # type: ignore[call-overload] else: raise ValueError( f"Expected a 1D array, got an array with shape {value.shape}" diff --git a/pandas/core/internals/concat.py b/pandas/core/internals/concat.py index 6b41d7a26080d..b34d3590f6a71 100644 --- a/pandas/core/internals/concat.py +++ b/pandas/core/internals/concat.py @@ -538,10 +538,10 @@ def _concatenate_join_units( # concatting with at least one EA means we are concatting a single column # the non-EA values are 2D arrays with shape (1, n) - # error: Invalid index type "Tuple[int, slice]" for - # "Union[ExtensionArray, ndarray]"; expected type "Union[int, slice, ndarray]" + # error: No overload variant of "__getitem__" of "ExtensionArray" matches + # argument type "Tuple[int, slice]" to_concat = [ - t if is_1d_only_ea_obj(t) else t[0, :] # type: ignore[index] + t if is_1d_only_ea_obj(t) else t[0, :] # type: ignore[call-overload] for t in to_concat ] concat_values = concat_compat(to_concat, axis=0, ea_compat_axis=True) diff --git a/pandas/core/internals/ops.py b/pandas/core/internals/ops.py index 5f03d6709dfa4..35caeea9b9067 100644 --- a/pandas/core/internals/ops.py +++ b/pandas/core/internals/ops.py @@ -106,28 +106,28 @@ def _get_same_shape_values( # TODO(EA2D): with 2D EAs only this first clause would be needed if not (left_ea or right_ea): - # error: Invalid index type "Tuple[Any, slice]" for "Union[ndarray, - # ExtensionArray]"; expected type "Union[int, slice, ndarray]" - lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[index] + # error: No overload variant of "__getitem__" of "ExtensionArray" matches + # argument type "Tuple[Union[ndarray, slice], slice]" + lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[call-overload] assert lvals.shape == rvals.shape, (lvals.shape, rvals.shape) elif left_ea and right_ea: assert lvals.shape == rvals.shape, (lvals.shape, rvals.shape) elif right_ea: # lvals are 2D, rvals are 1D - # error: Invalid index type "Tuple[Any, slice]" for "Union[ndarray, - # ExtensionArray]"; expected type "Union[int, slice, ndarray]" - lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[index] + # error: No overload variant of "__getitem__" of "ExtensionArray" matches + # argument type "Tuple[Union[ndarray, slice], slice]" + lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[call-overload] assert lvals.shape[0] == 1, lvals.shape - # error: Invalid index type "Tuple[int, slice]" for "Union[Any, - # ExtensionArray]"; expected type "Union[int, slice, ndarray]" - lvals = lvals[0, :] # type: ignore[index] + # error: No overload variant of "__getitem__" of "ExtensionArray" matches + # argument type "Tuple[int, slice]" + lvals = lvals[0, :] # type: ignore[call-overload] else: # lvals are 1D, rvals are 2D assert rvals.shape[0] == 1, rvals.shape - # error: Invalid index type "Tuple[int, slice]" for "Union[ndarray, - # ExtensionArray]"; expected type "Union[int, slice, ndarray]" - rvals = rvals[0, :] # type: ignore[index] + # error: No overload variant of "__getitem__" of "ExtensionArray" matches + # argument type "Tuple[int, slice]" + rvals = rvals[0, :] # type: ignore[call-overload] return lvals, rvals