Skip to content

Commit 4c6f81c

Browse files
cmp0xffcmp0xff
committed
feat: truediv for index[bool]
1 parent 5cba4c1 commit 4c6f81c

File tree

20 files changed

+290
-42
lines changed

20 files changed

+290
-42
lines changed

‎pandas-stubs/core/base.pyi‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,6 @@ NumListLike: TypeAlias = (
179179
| np_ndarray_complex
180180
| dict[str, np.ndarray]
181181
| Sequence[complex]
182-
| IndexOpsMixin[complex]
183182
)
184183

185184
@type_check_only

‎pandas-stubs/core/indexes/base.pyi‎

Lines changed: 158 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ from datetime import (
99
datetime,
1010
timedelta,
1111
)
12+
from pathlib import Path
1213
from typing import (
1314
Any,
1415
ClassVar,
@@ -42,11 +43,12 @@ from pandas import (
4243
from pandas.core.base import (
4344
ElementOpsMixin,
4445
IndexOpsMixin,
45-
NumListLike,
4646
Supports_ProtoAdd,
4747
Supports_ProtoMul,
4848
Supports_ProtoRAdd,
4949
Supports_ProtoRMul,
50+
Supports_ProtoRTrueDiv,
51+
Supports_ProtoTrueDiv,
5052
_ListLike,
5153
)
5254
from pandas.core.indexes.category import CategoricalIndex
@@ -674,7 +676,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
674676
@overload
675677
def __rsub__(self: Index[Never], other: DatetimeIndex) -> Never: ... # type: ignore[misc]
676678
@overload
677-
def __rsub__(self: Index[Never], other: complex | NumListLike | Index) -> Index: ...
679+
def __rsub__(self: Index[Never], other: complex | _ListLike | Index) -> Index: ...
678680
@overload
679681
def __rsub__(self, other: Index[Never]) -> Index: ...
680682
@overload
@@ -893,18 +895,166 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
893895
def __rmul__(
894896
self: Index[T_COMPLEX], other: np_ndarray_complex | Index[complex]
895897
) -> Index[complex]: ...
898+
@overload
899+
def __truediv__(
900+
self: Index[Never], other: complex | _ListLike | Index
901+
) -> Index: ...
902+
@overload
903+
def __truediv__(self, other: Index[Never]) -> Index: ...
904+
@overload
905+
def __truediv__(self: Index[bool], other: np_ndarray_bool) -> Never: ...
906+
@overload
907+
def __truediv__(
908+
self: Index[Timedelta],
909+
other: np_ndarray_bool | np_ndarray_complex | np_ndarray_dt,
910+
) -> Never: ...
911+
@overload
912+
def __truediv__(
913+
self: Supports_ProtoTrueDiv[_T_contra, S2],
914+
other: _T_contra | Sequence[_T_contra],
915+
) -> Index[S2]: ...
916+
@overload
917+
def __truediv__(
918+
self: Index[int],
919+
other: np_ndarray_bool | Index[bool],
920+
) -> Index[float]: ...
921+
@overload
922+
def __truediv__(
923+
self: Index[bool] | Index[int],
924+
other: Just[int] | Sequence[int] | np_ndarray_anyint | Index[int],
925+
) -> Index[float]: ...
926+
@overload
927+
def __truediv__(
928+
self: Index[float],
929+
other: np_ndarray_bool | np_ndarray_anyint | Index[bool] | Index[int],
930+
) -> Index[float]: ...
931+
@overload
932+
def __truediv__(
933+
self: Index[complex],
934+
other: np_ndarray_bool | np_ndarray_anyint | Index[bool] | Index[int],
935+
) -> Index[complex]: ...
936+
@overload
937+
def __truediv__(
938+
self: Index[bool] | Index[int],
939+
other: Just[float] | Sequence[Just[float]] | np_ndarray_float | Index[float],
940+
) -> Index[float]: ...
941+
@overload
942+
def __truediv__(
943+
self: Index[T_COMPLEX],
944+
other: Just[float] | Sequence[Just[float]] | np_ndarray_float | Index[float],
945+
) -> Index[T_COMPLEX]: ...
946+
@overload
947+
def __truediv__(
948+
self: Index[T_COMPLEX],
949+
other: (
950+
Just[complex]
951+
| Sequence[Just[complex]]
952+
| np_ndarray_complex
953+
| Index[complex]
954+
),
955+
) -> Index[complex]: ...
956+
@overload
957+
def __truediv__(
958+
self: Index[Timedelta],
959+
other: (
960+
Just[int]
961+
| Just[float]
962+
| Sequence[Just[int]]
963+
| Sequence[Just[float]]
964+
| np_ndarray_anyint
965+
| np_ndarray_float
966+
| Index[int]
967+
| Index[float]
968+
),
969+
) -> Index[Timedelta]: ...
970+
@overload
971+
def __truediv__(
972+
self: Index[Timedelta],
973+
other: np_ndarray_td | TimedeltaIndex | Index[Timedelta],
974+
) -> Index[float]: ...
975+
@overload
976+
def __truediv__(self, other: Path) -> Index: ...
977+
@overload
978+
def __rtruediv__(
979+
self: Index[Never], other: complex | _ListLike | Index
980+
) -> Index: ...
981+
@overload
982+
def __rtruediv__(self, other: Index[Never]) -> Index: ...
983+
@overload
984+
def __rtruediv__(self: Index[bool], other: np_ndarray_bool) -> Never: ...
985+
@overload
986+
def __rtruediv__(
987+
self: Index[Timedelta],
988+
other: np_ndarray_bool | np_ndarray_complex | np_ndarray_dt,
989+
) -> Never: ...
990+
@overload
991+
def __rtruediv__(
992+
self: Supports_ProtoRTrueDiv[_T_contra, S2],
993+
other: _T_contra | Sequence[_T_contra],
994+
) -> Index[S2]: ...
995+
@overload
996+
def __rtruediv__(
997+
self: Index[int], other: np_ndarray_bool | Index[bool]
998+
) -> Index[float]: ...
999+
@overload
1000+
def __rtruediv__(
1001+
self: Index[bool] | Index[int],
1002+
other: Just[int] | Sequence[int] | np_ndarray_anyint | Index[int],
1003+
) -> Index[float]: ...
1004+
@overload
1005+
def __rtruediv__( # type: ignore[misc]
1006+
self: Index[float],
1007+
other: np_ndarray_bool | np_ndarray_anyint | Index[bool] | Index[int],
1008+
) -> Index[float]: ...
1009+
@overload
1010+
def __rtruediv__(
1011+
self: Index[complex],
1012+
other: np_ndarray_bool | np_ndarray_anyint | Index[bool] | Index[int],
1013+
) -> Index[complex]: ...
1014+
@overload
1015+
def __rtruediv__(
1016+
self: Index[bool] | Index[int],
1017+
other: Just[float] | Sequence[Just[float]] | np_ndarray_float | Index[float],
1018+
) -> Index[float]: ...
1019+
@overload
1020+
def __rtruediv__(
1021+
self: Index[T_COMPLEX],
1022+
other: Just[float] | Sequence[Just[float]] | np_ndarray_float | Index[float],
1023+
) -> Index[T_COMPLEX]: ...
1024+
@overload
1025+
def __rtruediv__(
1026+
self: Index[T_COMPLEX],
1027+
other: (
1028+
Just[complex]
1029+
| Sequence[Just[complex]]
1030+
| np_ndarray_complex
1031+
| Index[complex]
1032+
),
1033+
) -> Index[complex]: ...
1034+
@overload
1035+
def __rtruediv__(
1036+
self: Index[Timedelta],
1037+
other: np_ndarray_td | TimedeltaIndex | Index[Timedelta],
1038+
) -> Index[float]: ...
1039+
@overload
1040+
def __rtruediv__(
1041+
self: Index[int] | Index[float],
1042+
other: (
1043+
timedelta
1044+
| Sequence[timedelta]
1045+
| np_ndarray_td
1046+
| TimedeltaIndex
1047+
| Index[Timedelta]
1048+
),
1049+
) -> Index[Timedelta]: ...
1050+
@overload
1051+
def __rtruediv__(self, other: Path) -> Index: ...
8961052
def __floordiv__(
8971053
self, other: float | Sequence[float] | Index[int] | Index[float]
8981054
) -> Self: ...
8991055
def __rfloordiv__(
9001056
self, other: float | Sequence[float] | Index[int] | Index[float]
9011057
) -> Self: ...
902-
def __truediv__(
903-
self, other: float | Sequence[float] | Index[int] | Index[float]
904-
) -> Self: ...
905-
def __rtruediv__(
906-
self, other: float | Sequence[float] | Index[int] | Index[float]
907-
) -> Self: ...
9081058
def infer_objects(self, copy: bool = True) -> Self: ...
9091059

9101060
@type_check_only

‎tests/indexes/arithmetic/bool/test_mul.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222
@pytest.fixture
2323
def left() -> "pd.Index[bool]":
24-
"""left operand"""
24+
"""Left operand"""
2525
lo = pd.Index([True, True, False])
2626
return check(assert_type(lo, "pd.Index[bool]"), pd.Index, np.bool_)
2727

‎tests/indexes/arithmetic/bool/test_truediv.py‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
import numpy as np
2+
from numpy import typing as npt # noqa: F401
3+
import pandas as pd
4+
import pytest
5+
from typing_extensions import (
6+
Never,
7+
assert_type,
8+
)
9+
10+
from tests import (
11+
TYPE_CHECKING_INVALID_USAGE,
12+
check,
13+
)
14+
15+
16+
@pytest.fixture
17+
def left() -> "pd.Index[bool]":
18+
"""Left operand"""
19+
return check(
20+
assert_type(pd.Index([True, False, True]), "pd.Index[bool]"), pd.Index, np.bool_
21+
)
22+
23+
24+
def test_truediv_py_scalar(left: "pd.Index[bool]") -> None:
25+
"""Test pd.Index[bool] / Python native scalars"""
26+
b, i, f, c = True, 1, 1.0, 1j
27+
28+
if TYPE_CHECKING_INVALID_USAGE:
29+
# TODO: python/mypy#20061
30+
_00 = left / b # pyright: ignore[reportOperatorIssue]
31+
check(assert_type(left / i, "pd.Index[float]"), pd.Index, np.floating)
32+
check(assert_type(left / f, "pd.Index[float]"), pd.Index, np.floating)
33+
check(assert_type(left / c, "pd.Index[complex]"), pd.Index, np.complexfloating)
34+
35+
if TYPE_CHECKING_INVALID_USAGE:
36+
# TODO: python/mypy#20061
37+
_10 = b / left # pyright: ignore[reportOperatorIssue]
38+
check(assert_type(i / left, "pd.Index[float]"), pd.Index, np.floating)
39+
check(assert_type(f / left, "pd.Index[float]"), pd.Index, np.floating)
40+
check(assert_type(c / left, "pd.Index[complex]"), pd.Index, np.complexfloating)
41+
42+
43+
def test_truediv_py_sequence(left: "pd.Index[bool]") -> None:
44+
"""Test pd.Index[bool] / Python native sequences"""
45+
b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j]
46+
47+
check(assert_type(left / b, "pd.Index[float]"), pd.Index, np.floating)
48+
check(assert_type(left / i, "pd.Index[float]"), pd.Index, np.floating)
49+
check(assert_type(left / f, "pd.Index[float]"), pd.Index, np.floating)
50+
check(assert_type(left / c, "pd.Index[complex]"), pd.Index, np.complexfloating)
51+
52+
check(assert_type(b / left, "pd.Index[float]"), pd.Index, np.floating)
53+
check(assert_type(i / left, "pd.Index[float]"), pd.Index, np.floating)
54+
check(assert_type(f / left, "pd.Index[float]"), pd.Index, np.floating)
55+
check(assert_type(c / left, "pd.Index[complex]"), pd.Index, np.complexfloating)
56+
57+
58+
def test_truediv_numpy_array(left: "pd.Index[bool]") -> None:
59+
"""Test pd.Index[bool] / numpy arrays"""
60+
b = np.array([True, False, True], np.bool_)
61+
i = np.array([2, 3, 5], np.int64)
62+
f = np.array([1.0, 2.0, 3.0], np.float64)
63+
c = np.array([1.1j, 2.2j, 4.1j], np.complex128)
64+
65+
if TYPE_CHECKING_INVALID_USAGE:
66+
assert_type(left / b, Never)
67+
check(assert_type(left / i, "pd.Index[float]"), pd.Index, np.floating)
68+
check(assert_type(left / f, "pd.Index[float]"), pd.Index, np.floating)
69+
check(assert_type(left / c, "pd.Index[complex]"), pd.Index, np.complexfloating)
70+
71+
# `numpy` typing gives the corresponding `ndarray`s in the static type
72+
# checking, where our `__rtruediv__` cannot override. At runtime, they return
73+
# `Index` with the correct element type.
74+
if TYPE_CHECKING_INVALID_USAGE:
75+
assert_type(b / left, "npt.NDArray[np.float64]")
76+
check(assert_type(i / left, "npt.NDArray[np.float64]"), pd.Index, np.floating)
77+
check(assert_type(f / left, "npt.NDArray[np.float64]"), pd.Index, np.floating)
78+
check(
79+
assert_type(c / left, "npt.NDArray[np.complex128]"),
80+
pd.Index,
81+
np.complexfloating,
82+
)
83+
84+
85+
def test_truediv_pd_index(left: "pd.Index[bool]") -> None:
86+
"""Test pd.Index[bool] / pandas Indexes"""
87+
b = pd.Index([True, False, True])
88+
i = pd.Index([2, 3, 5])
89+
f = pd.Index([1.0, 2.0, 3.0])
90+
c = pd.Index([1.1j, 2.2j, 4.1j])
91+
92+
if TYPE_CHECKING_INVALID_USAGE:
93+
_00 = left / b # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
94+
check(assert_type(left / i, "pd.Index[float]"), pd.Index, np.floating)
95+
check(assert_type(left / f, "pd.Index[float]"), pd.Index, np.floating)
96+
check(assert_type(left / c, "pd.Index[complex]"), pd.Index, np.complexfloating)
97+
98+
if TYPE_CHECKING_INVALID_USAGE:
99+
_10 = b / left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
100+
check(assert_type(i / left, "pd.Index[float]"), pd.Index, np.floating)
101+
check(assert_type(f / left, "pd.Index[float]"), pd.Index, np.floating)
102+
check(assert_type(c / left, "pd.Index[complex]"), pd.Index, np.complexfloating)

‎tests/indexes/arithmetic/complex/test_mul.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222
@pytest.fixture
2323
def left() -> "pd.Index[complex]":
24-
"""left operand"""
24+
"""Left operand"""
2525
lo = pd.Index([1j, 2j, 3j])
2626
return check(assert_type(lo, "pd.Index[complex]"), pd.Index, np.complexfloating)
2727

‎tests/indexes/arithmetic/float/test_mul.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222
@pytest.fixture
2323
def left() -> "pd.Index[float]":
24-
"""left operand"""
24+
"""Left operand"""
2525
lo = pd.Index([1.0, 2.0, 3.0])
2626
return check(assert_type(lo, "pd.Index[float]"), pd.Index, np.floating)
2727

‎tests/indexes/arithmetic/int/test_mul.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222
@pytest.fixture
2323
def left() -> "pd.Index[int]":
24-
"""left operand"""
24+
"""Left operand"""
2525
lo = pd.Index([1, 2, 3])
2626
return check(assert_type(lo, "pd.Index[int]"), pd.Index, np.integer)
2727

‎tests/indexes/arithmetic/str/test_mul.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222
@pytest.fixture
2323
def left() -> "pd.Index[str]":
24-
"""left operand"""
24+
"""Left operand"""
2525
lo = pd.Index(["1", "2", "3"])
2626
return check(assert_type(lo, "pd.Index[str]"), pd.Index, str)
2727

‎tests/indexes/arithmetic/test_mul.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313

1414
@pytest.fixture
1515
def left_i() -> pd.Index:
16-
"""left operand"""
16+
"""Left operand"""
1717
lo = pd.MultiIndex.from_tuples([(1,), (2,), (3,)]).levels[0]
1818
return check(assert_type(lo, pd.Index), pd.Index)
1919

‎tests/indexes/arithmetic/timedelta/test_mul.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
@pytest.fixture
1919
def left() -> "pd.Index[pd.Timedelta]":
20-
"""left operand"""
20+
"""Left operand"""
2121
# pandas-dev/pandas#62524: An index of Python native timedeltas can be
2222
# produced, instead of a TimedeltaIndex, hence this test file
2323
lo = pd.Index([1]) * [timedelta(seconds=1)] # left operand

0 commit comments

Comments
 (0)