diff --git a/pandas-stubs/core/arraylike.pyi b/pandas-stubs/core/arraylike.pyi index 51182a976..df9f8bd8c 100644 --- a/pandas-stubs/core/arraylike.pyi +++ b/pandas-stubs/core/arraylike.pyi @@ -1,46 +1,40 @@ from typing import ( Any, - Protocol, + TypeVar, ) -from pandas import DataFrame - -class OpsMixinProtocol(Protocol): ... +_OpsMixinT = TypeVar("_OpsMixinT", bound=OpsMixin) class OpsMixin: - def __eq__(self: OpsMixinProtocol, other: object) -> DataFrame: ... # type: ignore[override] - def __ne__(self: OpsMixinProtocol, other: object) -> DataFrame: ... # type: ignore[override] - def __lt__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __le__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __gt__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __ge__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... + def __eq__(self: _OpsMixinT, other: object) -> _OpsMixinT: ... # type: ignore[override] + def __ne__(self: _OpsMixinT, other: object) -> _OpsMixinT: ... # type: ignore[override] + def __lt__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __le__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __gt__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __ge__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... # ------------------------------------------------------------- # Logical Methods - def __and__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __rand__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __or__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __ror__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __xor__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __rxor__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... + def __and__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __rand__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __or__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __ror__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __xor__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __rxor__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... # ------------------------------------------------------------- # Arithmetic Methods - def __add__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __radd__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __sub__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __rsub__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __mul__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __rmul__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __truediv__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __rtruediv__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __floordiv__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __rfloordiv__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __mod__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __rmod__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __divmod__( - self: OpsMixinProtocol, other: DataFrame - ) -> tuple[DataFrame, DataFrame]: ... - def __rdivmod__( - self: OpsMixinProtocol, other: DataFrame - ) -> tuple[DataFrame, DataFrame]: ... - def __pow__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... - def __rpow__(self: OpsMixinProtocol, other: Any) -> DataFrame: ... + def __add__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __radd__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __sub__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __rsub__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __mul__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __rmul__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __truediv__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __rtruediv__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __floordiv__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __rfloordiv__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __mod__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __rmod__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __divmod__(self: _OpsMixinT, other: Any) -> tuple[_OpsMixinT, _OpsMixinT]: ... + def __rdivmod__(self: _OpsMixinT, other: Any) -> tuple[_OpsMixinT, _OpsMixinT]: ... + def __pow__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... + def __rpow__(self: _OpsMixinT, other: Any) -> _OpsMixinT: ... diff --git a/pandas-stubs/core/base.pyi b/pandas-stubs/core/base.pyi index 213cc5d93..1d1902ae8 100644 --- a/pandas-stubs/core/base.pyi +++ b/pandas-stubs/core/base.pyi @@ -6,6 +6,7 @@ from typing import ( import numpy as np from pandas import Index from pandas.core.accessor import DirNamesMixin +from pandas.core.arraylike import OpsMixin from pandas.core.arrays import ExtensionArray from pandas.core.arrays.categorical import Categorical @@ -27,7 +28,7 @@ class SelectionMixin(Generic[NDFrameT]): def ndim(self) -> int: ... def __getitem__(self, key): ... -class IndexOpsMixin: +class IndexOpsMixin(OpsMixin): __array_priority__: int = ... def transpose(self, *args, **kwargs) -> IndexOpsMixin: ... @property diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index b0507fa1a..98a6b72bd 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -24,6 +24,7 @@ from pandas.core.base import ( ) from pandas.core.indexes.numeric import NumericIndex from pandas.core.strings import StringMethods +from typing_extensions import Never from pandas._typing import ( T1, @@ -147,14 +148,12 @@ class Index(IndexOpsMixin, PandasObject): def duplicated( self, keep: Literal["first", "last", False] = ... ) -> np_ndarray_bool: ... - def __add__(self, other) -> Index: ... - def __radd__(self, other) -> Index: ... - def __iadd__(self, other) -> Index: ... - def __sub__(self, other) -> Index: ... - def __rsub__(self, other) -> Index: ... - def __and__(self, other) -> Index: ... - def __or__(self, other) -> Index: ... - def __xor__(self, other) -> Index: ... + def __and__(self, other: Never) -> Never: ... + def __rand__(self, other: Never) -> Never: ... + def __or__(self, other: Never) -> Never: ... + def __ror__(self, other: Never) -> Never: ... + def __xor__(self, other: Never) -> Never: ... + def __rxor__(self, other: Never) -> Never: ... def __neg__(self: IndexT) -> IndexT: ... def __nonzero__(self) -> None: ... __bool__ = ... @@ -226,10 +225,10 @@ class Index(IndexOpsMixin, PandasObject): def __eq__(self, other: object) -> np_ndarray_bool: ... # type: ignore[override] def __iter__(self) -> Iterator[IndexIterScalar | tuple[Hashable, ...]]: ... def __ne__(self, other: object) -> np_ndarray_bool: ... # type: ignore[override] - def __le__(self, other: Index | Scalar) -> np_ndarray_bool: ... - def __ge__(self, other: Index | Scalar) -> np_ndarray_bool: ... - def __lt__(self, other: Index | Scalar) -> np_ndarray_bool: ... - def __gt__(self, other: Index | Scalar) -> np_ndarray_bool: ... + def __le__(self, other: Index | Scalar) -> np_ndarray_bool: ... # type: ignore[override] + def __ge__(self, other: Index | Scalar) -> np_ndarray_bool: ... # type: ignore[override] + def __lt__(self, other: Index | Scalar) -> np_ndarray_bool: ... # type: ignore[override] + def __gt__(self, other: Index | Scalar) -> np_ndarray_bool: ... # type: ignore[override] def ensure_index_from_sequences( sequences: Sequence[Sequence[Dtype]], names: list[str] = ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index c9075dcb5..a4bd5c0ef 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1277,7 +1277,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): self, other: num | _str | Timedelta | _ListLike | Series[S1] | np.timedelta64 ) -> Series: ... # ignore needed for mypy as we want different results based on the arguments - @overload + @overload # type: ignore[override] def __and__( # type: ignore[misc] self, other: bool | list[bool] | np_ndarray_bool | Series[bool] ) -> Series[bool]: ... @@ -1288,29 +1288,17 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): # def __array__(self, dtype: Optional[_bool] = ...) -> _np_ndarray def __div__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... def __eq__(self, other: object) -> Series[_bool]: ... # type: ignore[override] - def __floordiv__(self, other: num | _ListLike | Series[S1]) -> Series[int]: ... - def __ge__( + def __floordiv__(self, other: num | _ListLike | Series[S1]) -> Series[int]: ... # type: ignore[override] + def __ge__( # type: ignore[override] self, other: S1 | _ListLike | Series[S1] | datetime | timedelta ) -> Series[_bool]: ... - def __gt__( + def __gt__( # type: ignore[override] self, other: S1 | _ListLike | Series[S1] | datetime | timedelta ) -> Series[_bool]: ... - # def __iadd__(self, other: S1) -> Series[S1]: ... - # def __iand__(self, other: S1) -> Series[_bool]: ... - # def __idiv__(self, other: S1) -> Series[S1]: ... - # def __ifloordiv__(self, other: S1) -> Series[S1]: ... - # def __imod__(self, other: S1) -> Series[S1]: ... - # def __imul__(self, other: S1) -> Series[S1]: ... - # def __ior__(self, other: S1) -> Series[_bool]: ... - # def __ipow__(self, other: S1) -> Series[S1]: ... - # def __isub__(self, other: S1) -> Series[S1]: ... - # def __itruediv__(self, other: S1) -> Series[S1]: ... - # def __itruediv__(self, other) -> None: ... - # def __ixor__(self, other: S1) -> Series[_bool]: ... - def __le__( + def __le__( # type: ignore[override] self, other: S1 | _ListLike | Series[S1] | datetime | timedelta ) -> Series[_bool]: ... - def __lt__( + def __lt__( # type: ignore[override] self, other: S1 | _ListLike | Series[S1] | datetime | timedelta ) -> Series[_bool]: ... @overload @@ -1323,7 +1311,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def __ne__(self, other: object) -> Series[_bool]: ... # type: ignore[override] def __pow__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... # ignore needed for mypy as we want different results based on the arguments - @overload + @overload # type: ignore[override] def __or__( # type: ignore[misc] self, other: bool | list[bool] | np_ndarray_bool | Series[bool] ) -> Series[bool]: ... @@ -1333,7 +1321,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): ) -> Series[int]: ... def __radd__(self, other: num | _str | _ListLike | Series[S1]) -> Series[S1]: ... # ignore needed for mypy as we want different results based on the arguments - @overload + @overload # type: ignore[override] def __rand__( # type: ignore[misc] self, other: bool | list[bool] | np_ndarray_bool | Series[bool] ) -> Series[bool]: ... @@ -1342,14 +1330,14 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): self, other: int | list[int] | np_ndarray_anyint | Series[int] ) -> Series[int]: ... def __rdiv__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... - def __rdivmod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... + def __rdivmod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... # type: ignore[override] def __rfloordiv__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... def __rmod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... def __rmul__(self, other: num | _ListLike | Series) -> Series: ... def __rnatmul__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... def __rpow__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... # ignore needed for mypy as we want different results based on the arguments - @overload + @overload # type: ignore[override] def __ror__( # type: ignore[misc] self, other: bool | list[bool] | np_ndarray_bool | Series[bool] ) -> Series[bool]: ... @@ -1363,7 +1351,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def __rtruediv__(self, other: num | _ListLike | Series[S1]) -> Series: ... # ignore needed for mypy as we want different results based on the arguments - @overload + @overload # type: ignore[override] def __rxor__( # type: ignore[misc] self, other: bool | list[bool] | np_ndarray_bool | Series[bool] ) -> Series[bool]: ... @@ -1389,7 +1377,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def __sub__(self, other: num | _ListLike | Series) -> Series: ... def __truediv__(self, other: num | _ListLike | Series[S1]) -> Series: ... # ignore needed for mypy as we want different results based on the arguments - @overload + @overload # type: ignore[override] def __xor__( # type: ignore[misc] self, other: bool | list[bool] | np_ndarray_bool | Series[bool] ) -> Series[bool]: ... diff --git a/pyproject.toml b/pyproject.toml index e823708d5..61f96bd3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ poethepoet = ">=0.16.5" loguru = ">=0.6.0" pandas = "1.5.3" numpy = ">=1.24.1" -typing-extensions = ">=4.2.0" +typing-extensions = ">=4.4.0" matplotlib = ">=3.5.1" pre-commit = ">=2.19.0" black = ">=22.12.0" diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 84867d8aa..e4931a9e6 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -13,7 +13,10 @@ from numpy import typing as npt import pandas as pd from pandas.core.indexes.numeric import NumericIndex -from typing_extensions import assert_type +from typing_extensions import ( + Never, + assert_type, +) from pandas._typing import Scalar @@ -21,6 +24,7 @@ from pandas._typing import IndexIterScalar from tests import ( + TYPE_CHECKING_INVALID_USAGE, check, pytest_warns_bounded, ) @@ -684,3 +688,72 @@ def test_sorted_and_list() -> None: ), list, ) + + +def test_index_operators() -> None: + # GH 405 + i1 = pd.Index([1, 2, 3]) + i2 = pd.Index([4, 5, 6]) + + check(assert_type(i1 + i2, pd.Index), pd.Index) + check(assert_type(i1 + 10, pd.Index), pd.Index) + check(assert_type(10 + i1, pd.Index), pd.Index) + check(assert_type(i1 - i2, pd.Index), pd.Index) + check(assert_type(i1 - 10, pd.Index), pd.Index) + check(assert_type(10 - i1, pd.Index), pd.Index) + check(assert_type(i1 * i2, pd.Index), pd.Index) + check(assert_type(i1 * 10, pd.Index), pd.Index) + check(assert_type(10 * i1, pd.Index), pd.Index) + check(assert_type(i1 / i2, pd.Index), pd.Index) + check(assert_type(i1 / 10, pd.Index), pd.Index) + check(assert_type(10 / i1, pd.Index), pd.Index) + check(assert_type(i1 // i2, pd.Index), pd.Index) + check(assert_type(i1 // 10, pd.Index), pd.Index) + check(assert_type(10 // i1, pd.Index), pd.Index) + check(assert_type(i1**i2, pd.Index), pd.Index) + check(assert_type(i1**2, pd.Index), pd.Index) + check(assert_type(2**i1, pd.Index), pd.Index) + check(assert_type(i1 % i2, pd.Index), pd.Index) + check(assert_type(i1 % 10, pd.Index), pd.Index) + check(assert_type(10 % i1, pd.Index), pd.Index) + check(assert_type(divmod(i1, i2), Tuple[pd.Index, pd.Index]), tuple) + check(assert_type(divmod(i1, 10), Tuple[pd.Index, pd.Index]), tuple) + check(assert_type(divmod(10, i1), Tuple[pd.Index, pd.Index]), tuple) + + if TYPE_CHECKING_INVALID_USAGE: + assert_type( + i1 & i2, # type:ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + Never, + ) + assert_type( # type: ignore[assert-type] + i1 & 10, # type:ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + Never, + ) + assert_type( # type: ignore[assert-type] + 10 & i1, # type:ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + Never, + ) + assert_type( + i1 | i2, # type:ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + Never, + ) + assert_type( # type: ignore[assert-type] + i1 | 10, # type:ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + Never, + ) + assert_type( # type: ignore[assert-type] + 10 | i1, # type:ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + Never, + ) + assert_type( + i1 ^ i2, # type:ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + Never, + ) + assert_type( # type: ignore[assert-type] + i1 ^ 10, # type:ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + Never, + ) + assert_type( # type: ignore[assert-type] + 10 ^ i1, # type:ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + Never, + ) diff --git a/tests/test_series.py b/tests/test_series.py index e2f5c83d7..2d8adf578 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -12,6 +12,7 @@ Iterator, List, Sequence, + Tuple, TypeVar, cast, ) @@ -452,6 +453,8 @@ def test_types_element_wise_arithmetic() -> None: res_pow: pd.Series = s ** s2.abs() res_pow2: pd.Series = s.pow(s2.abs(), fill_value=0) + check(assert_type(divmod(s, s2), Tuple[pd.Series, pd.Series]), tuple) + def test_types_scalar_arithmetic() -> None: s = pd.Series([0, 1, -10])