Skip to content

add operators for Index #504

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.

Sign up for GitHub

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 4 commits into from
Feb 3, 2023
Merged

add operators for Index #504

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2d0f53f
add operators for Index
Dr-Irv Jan 9, 2023
0a66721
Fix OpsMixin using TypeVar
Dr-Irv Jan 10, 2023
3b3f280
detect Never as argument
Dr-Irv Jan 10, 2023
3d99b2b
merge in main. resolve conflicts
Dr-Irv Jan 31, 2023
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
66 changes: 30 additions & 36 deletions pandas-stubs/core/arraylike.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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: ...
3 changes: 2 additions & 1 deletion pandas-stubs/core/base.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand Down
23 changes: 11 additions & 12 deletions pandas-stubs/core/indexes/base.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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__ = ...
Expand Down Expand Up @@ -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] = ...
Expand Down
36 changes: 12 additions & 24 deletions pandas-stubs/core/series.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]: ...
Expand All @@ -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
Expand All @@ -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]: ...
Expand All @@ -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]: ...
Expand All @@ -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]: ...
Expand All @@ -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]: ...
Expand All @@ -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]: ...
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
75 changes: 74 additions & 1 deletion tests/test_indexes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,18 @@
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

if TYPE_CHECKING:
from pandas._typing import IndexIterScalar

from tests import (
TYPE_CHECKING_INVALID_USAGE,
check,
pytest_warns_bounded,
)
Expand Down Expand Up @@ -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)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not needed but could also check the tuple elements: check(..., tuple, pd.Index)

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]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some of the checks have an ignore for assert_type and some don't. Why do some need an ignore?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It has to do with how mypy processes Never with certain operators and arguments. If you look carefully, you'll see that the # ignore[assert-type] is occurring on any operator that has a numerical operand. Without the ignore, mypy says the assert_type() is wrong, thinking the result is Any. Probably some kind of mypy error, although the interpretation of the Never result is ambiguous.

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,
)
3 changes: 3 additions & 0 deletions tests/test_series.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
Iterator,
List,
Sequence,
Tuple,
TypeVar,
cast,
)
Expand Down Expand Up @@ -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])
Expand Down