Skip to content

Commit f15f4ba

Browse files
cmp0xffcmp0xff
committed
floordiv
1 parent ce229fc commit f15f4ba

File tree

5 files changed

+569
-2
lines changed

5 files changed

+569
-2
lines changed

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

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1026,7 +1026,23 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
10261026
@overload
10271027
def __rtruediv__(self, other: Path) -> Index: ...
10281028
@overload
1029-
def __floordiv__(self, other: Index[Never]) -> Index: ...
1029+
def __floordiv__(
1030+
self: Index[Never],
1031+
other: (
1032+
float
1033+
| Sequence[float]
1034+
| np_ndarray_bool
1035+
| np_ndarray_anyint
1036+
| np_ndarray_float
1037+
| Index[bool]
1038+
| Index[int]
1039+
| Index[float]
1040+
),
1041+
) -> Index: ...
1042+
@overload
1043+
def __floordiv__(
1044+
self: Index[bool] | Index[int] | Index[float], other: Index[Never]
1045+
) -> Index: ...
10301046
@overload
10311047
def __floordiv__(
10321048
self: Index[int] | Index[float],
@@ -1063,7 +1079,21 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
10631079
other: float | Sequence[float] | np_ndarray_float | Index[float],
10641080
) -> Index[float]: ...
10651081
@overload
1066-
def __rfloordiv__(self, other: Index[Never]) -> Index: ... # type: ignore[overload-overlap]
1082+
def __rfloordiv__(
1083+
self: Index[Never],
1084+
other: (
1085+
float
1086+
| Sequence[float]
1087+
| np_ndarray_bool
1088+
| np_ndarray_anyint
1089+
| np_ndarray_float
1090+
| Index[bool]
1091+
| Index[int]
1092+
| Index[float]
1093+
),
1094+
) -> Index: ...
1095+
@overload
1096+
def __rfloordiv__(self: Index[bool] | Index[int] | Index[float], other: Index[Never]) -> Index: ... # type: ignore[overload-overlap]
10671097
@overload
10681098
def __rfloordiv__(
10691099
self: Index[int] | Index[float],

‎pandas-stubs/core/series.pyi‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2151,6 +2151,19 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
21512151
def __and__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ...
21522152
def __eq__(self, other: object) -> Series[_bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
21532153
@overload
2154+
def __floordiv__(
2155+
self: Series[Never],
2156+
other: (
2157+
float
2158+
| Sequence[float]
2159+
| np_ndarray_bool
2160+
| np_ndarray_anyint
2161+
| np_ndarray_float
2162+
| Index
2163+
| Series
2164+
),
2165+
) -> Series: ...
2166+
@overload
21542167
def __floordiv__(self, other: Index[Never] | Series[Never]) -> Series: ...
21552168
@overload
21562169
def __floordiv__(
@@ -2300,6 +2313,19 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
23002313
axis: AxisIndex | None = 0,
23012314
) -> Series[int]: ...
23022315
@overload
2316+
def __rfloordiv__(
2317+
self: Series[Never],
2318+
other: (
2319+
float
2320+
| Sequence[float]
2321+
| np_ndarray_bool
2322+
| np_ndarray_anyint
2323+
| np_ndarray_float
2324+
| Index
2325+
| Series
2326+
),
2327+
) -> Series: ...
2328+
@overload
23032329
def __rfloordiv__(self, other: Index[Never] | Series[Never]) -> Series: ...
23042330
@overload
23052331
def __rfloordiv__(

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

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
from datetime import (
2+
datetime,
3+
timedelta,
4+
)
5+
from typing import Any
6+
7+
import numpy as np
8+
import pandas as pd
9+
import pytest
10+
from typing_extensions import (
11+
Never,
12+
assert_type,
13+
)
14+
15+
from tests import (
16+
TYPE_CHECKING_INVALID_USAGE,
17+
check,
18+
)
19+
20+
21+
@pytest.fixture
22+
def left() -> "pd.Index[float]":
23+
"""Left operand"""
24+
lo = pd.Index([1.2, 2.4, 3.6])
25+
return check(assert_type(lo, "pd.Index[float]"), pd.Index, np.floating)
26+
27+
28+
def test_floordiv_py_scalar(left: "pd.Index[float]") -> None:
29+
"""Test pd.Index[float] // Python native scalars"""
30+
b, i, f, c = True, 1, 1.0, 1j
31+
s, d = datetime(2025, 9, 27), timedelta(seconds=1)
32+
33+
check(assert_type(left // b, "pd.Index[float]"), pd.Index, np.floating)
34+
check(assert_type(left // i, "pd.Index[float]"), pd.Index, np.floating)
35+
check(assert_type(left // f, "pd.Index[float]"), pd.Index, np.floating)
36+
if TYPE_CHECKING_INVALID_USAGE:
37+
_03 = left // c # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
38+
_04 = left // s # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
39+
_05 = left // d # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
40+
41+
check(assert_type(b // left, "pd.Index[float]"), pd.Index, np.floating)
42+
check(assert_type(i // left, "pd.Index[float]"), pd.Index, np.floating)
43+
check(assert_type(f // left, "pd.Index[float]"), pd.Index, np.floating)
44+
if TYPE_CHECKING_INVALID_USAGE:
45+
_13 = c // left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
46+
_14 = s // left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
47+
check(assert_type(d // left, pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta)
48+
49+
50+
def test_floordiv_py_sequence(left: "pd.Index[float]") -> None:
51+
"""Test pd.Index[float] // Python native sequences"""
52+
b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j]
53+
s = [datetime(2025, 10, 27 + d) for d in range(3)]
54+
d = [timedelta(seconds=s) for s in range(3)]
55+
56+
check(assert_type(left // b, "pd.Index[float]"), pd.Index, np.floating)
57+
check(assert_type(left // i, "pd.Index[float]"), pd.Index, np.floating)
58+
check(assert_type(left // f, "pd.Index[float]"), pd.Index, np.floating)
59+
if TYPE_CHECKING_INVALID_USAGE:
60+
_03 = left // c # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
61+
_04 = left // s # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
62+
_05 = left // d # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
63+
64+
check(assert_type(b // left, "pd.Index[float]"), pd.Index, np.floating)
65+
check(assert_type(i // left, "pd.Index[float]"), pd.Index, np.floating)
66+
check(assert_type(f // left, "pd.Index[float]"), pd.Index, np.floating)
67+
if TYPE_CHECKING_INVALID_USAGE:
68+
_13 = c // left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
69+
_14 = s // left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
70+
_15 = d // left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
71+
72+
73+
def test_floordiv_numpy_array(left: "pd.Index[float]") -> None:
74+
"""Test pd.Index[float] // numpy arrays"""
75+
b = np.array([True, False, True], np.bool_)
76+
i = np.array([2, 3, 5], np.int64)
77+
f = np.array([1.0, 2.0, 3.0], np.float64)
78+
c = np.array([1.1j, 2.2j, 4.1j], np.complex128)
79+
s = np.array(
80+
[np.datetime64(f"2025-10-{d:02d}") for d in (23, 24, 25)], np.datetime64
81+
)
82+
d = np.array([np.timedelta64(s + 1, "s") for s in range(3)], np.timedelta64)
83+
84+
check(assert_type(left // b, "pd.Index[float]"), pd.Index, np.floating)
85+
check(assert_type(left // i, "pd.Index[float]"), pd.Index, np.floating)
86+
check(assert_type(left // f, "pd.Index[float]"), pd.Index, np.floating)
87+
if TYPE_CHECKING_INVALID_USAGE:
88+
assert_type(left // c, Never)
89+
assert_type(left // s, Never)
90+
assert_type(left // d, Never)
91+
92+
# `numpy` typing gives the corresponding `ndarray`s in the static type
93+
# checking, where our `__rfloordiv__` cannot override. At runtime, they lead to
94+
# errors or pd.Index.
95+
check(b // left, pd.Index, np.floating)
96+
check(i // left, pd.Index, np.floating)
97+
check(f // left, pd.Index, np.floating)
98+
if TYPE_CHECKING_INVALID_USAGE:
99+
assert_type(c // left, Any)
100+
assert_type(s // left, Any)
101+
check(
102+
assert_type(d // left, "np.typing.NDArray[np.int64]"),
103+
pd.TimedeltaIndex,
104+
pd.Timedelta,
105+
)
106+
107+
108+
def test_floordiv_pd_index(left: "pd.Index[float]") -> None:
109+
"""Test pd.Index[float] // pandas Indexes"""
110+
b = pd.Index([True, False, True])
111+
i = pd.Index([2, 3, 5])
112+
f = pd.Index([1.0, 2.0, 3.0])
113+
c = pd.Index([1.1j, 2.2j, 4.1j])
114+
s = pd.Index([datetime(2025, 10, d) for d in (27, 28, 29)])
115+
d = pd.Index([timedelta(seconds=s + 1) for s in range(3)])
116+
117+
check(assert_type(left // b, "pd.Index[float]"), pd.Index, np.floating)
118+
check(assert_type(left // i, "pd.Index[float]"), pd.Index, np.floating)
119+
check(assert_type(left // f, "pd.Index[float]"), pd.Index, np.floating)
120+
if TYPE_CHECKING_INVALID_USAGE:
121+
_03 = left // c # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
122+
_04 = left // s # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
123+
_05 = left // d # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
124+
125+
check(assert_type(b // left, "pd.Index[float]"), pd.Index, np.floating)
126+
check(assert_type(i // left, "pd.Index[float]"), pd.Index, np.floating)
127+
check(assert_type(f // left, "pd.Index[float]"), pd.Index, np.floating)
128+
if TYPE_CHECKING_INVALID_USAGE:
129+
_13 = c // left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
130+
_14 = s // left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
131+
check(assert_type(d // left, pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta)

‎tests/indexes/arithmetic/test_floordiv.py‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
from datetime import (
2+
datetime,
3+
timedelta,
4+
)
5+
from typing import Any
6+
7+
import numpy as np
8+
import pandas as pd
9+
import pytest
10+
from typing_extensions import (
11+
Never,
12+
assert_type,
13+
)
14+
15+
from tests import (
16+
TYPE_CHECKING_INVALID_USAGE,
17+
check,
18+
)
19+
20+
21+
@pytest.fixture
22+
def left_i() -> pd.Index:
23+
"""Left operand"""
24+
lo = pd.MultiIndex.from_arrays([[1, 2, 3]]).levels[0]
25+
return check(assert_type(lo, pd.Index), pd.Index, np.integer)
26+
27+
28+
def test_floordiv_py_scalar(left_i: pd.Index) -> None:
29+
"""Test pd.Index[int] // Python native scalars"""
30+
b, i, f, c = True, 1, 1.0, 1j
31+
s, d = datetime(2025, 9, 27), timedelta(seconds=1)
32+
33+
check(assert_type(left_i // b, pd.Index), pd.Index, np.integer)
34+
check(assert_type(left_i // i, pd.Index), pd.Index, np.integer)
35+
check(assert_type(left_i // f, pd.Index), pd.Index, np.floating)
36+
if TYPE_CHECKING_INVALID_USAGE:
37+
_03 = left_i // c # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
38+
_04 = left_i // s # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
39+
_05 = left_i // d # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
40+
41+
check(assert_type(b // left_i, pd.Index), pd.Index, np.integer)
42+
check(assert_type(i // left_i, pd.Index), pd.Index, np.integer)
43+
check(assert_type(f // left_i, pd.Index), pd.Index, np.floating)
44+
if TYPE_CHECKING_INVALID_USAGE:
45+
_13 = c // left_i # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
46+
_14 = s // left_i # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
47+
check(assert_type(d // left_i, pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta)
48+
49+
50+
def test_floordiv_py_sequence(left_i: pd.Index) -> None:
51+
"""Test pd.Index[int] // Python native sequences"""
52+
b, i, f, c = [True, True, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j]
53+
s = [datetime(2025, 10, 27 + d) for d in range(3)]
54+
d = [timedelta(seconds=s) for s in range(3)]
55+
56+
check(assert_type(left_i // b, pd.Index), pd.Index, np.integer)
57+
check(assert_type(left_i // i, pd.Index), pd.Index, np.integer)
58+
check(assert_type(left_i // f, pd.Index), pd.Index, np.floating)
59+
if TYPE_CHECKING_INVALID_USAGE:
60+
_03 = left_i // c # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
61+
_04 = left_i // s # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
62+
_05 = left_i // d # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
63+
64+
check(assert_type(b // left_i, pd.Index), pd.Index, np.integer)
65+
check(assert_type(i // left_i, pd.Index), pd.Index, np.integer)
66+
check(assert_type(f // left_i, pd.Index), pd.Index, np.floating)
67+
if TYPE_CHECKING_INVALID_USAGE:
68+
_13 = c // left_i # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
69+
_14 = s // left_i # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
70+
_15 = d // left_i # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
71+
72+
73+
def test_floordiv_numpy_array(left_i: pd.Index) -> None:
74+
"""Test pd.Index[int] // numpy arrays"""
75+
b = np.array([True, True, True], np.bool_)
76+
i = np.array([2, 3, 5], np.int64)
77+
f = np.array([1.0, 2.0, 3.0], np.float64)
78+
c = np.array([1.1j, 2.2j, 4.1j], np.complex128)
79+
s = np.array(
80+
[np.datetime64(f"2025-10-{d:02d}") for d in (23, 24, 25)], np.datetime64
81+
)
82+
d = np.array([np.timedelta64(s + 1, "s") for s in range(3)], np.timedelta64)
83+
84+
check(assert_type(left_i // b, pd.Index), pd.Index, np.integer)
85+
check(assert_type(left_i // i, pd.Index), pd.Index, np.integer)
86+
check(assert_type(left_i // f, pd.Index), pd.Index, np.floating)
87+
if TYPE_CHECKING_INVALID_USAGE:
88+
assert_type(left_i // c, Never)
89+
assert_type(left_i // s, Never)
90+
assert_type(left_i // d, Never)
91+
92+
# `numpy` typing gives the corresponding `ndarray`s in the static type
93+
# checking, where our `__rfloordiv__` cannot override. At runtime, they lead to
94+
# errors or pd.Index.
95+
check(b // left_i, pd.Index, np.integer)
96+
check(i // left_i, pd.Index, np.integer)
97+
check(f // left_i, pd.Index, np.floating)
98+
if TYPE_CHECKING_INVALID_USAGE:
99+
assert_type(c // left_i, Any)
100+
assert_type(s // left_i, Any)
101+
assert_type(d // left_i, "np.typing.NDArray[np.int64]")
102+
103+
104+
def test_floordiv_pd_index(left_i: pd.Index) -> None:
105+
"""Test pd.Index[int] // pandas Indexes"""
106+
b = pd.Index([True, True, True])
107+
i = pd.Index([2, 3, 5])
108+
f = pd.Index([1.0, 2.0, 3.0])
109+
c = pd.Index([1.1j, 2.2j, 4.1j])
110+
s = pd.Index([datetime(2025, 10, d) for d in (27, 28, 29)])
111+
d = pd.Index([timedelta(seconds=s + 1) for s in range(3)])
112+
113+
check(assert_type(left_i // b, pd.Index), pd.Index, np.integer)
114+
check(assert_type(left_i // i, pd.Index), pd.Index, np.integer)
115+
check(assert_type(left_i // f, pd.Index), pd.Index, np.floating)
116+
if TYPE_CHECKING_INVALID_USAGE:
117+
_03 = left_i // c # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
118+
_04 = left_i // s # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
119+
_05 = left_i // d # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
120+
121+
check(assert_type(b // left_i, pd.Index), pd.Index, np.integer)
122+
check(assert_type(i // left_i, pd.Index), pd.Index, np.integer)
123+
check(assert_type(f // left_i, pd.Index), pd.Index, np.floating)
124+
if TYPE_CHECKING_INVALID_USAGE:
125+
_13 = c // left_i # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
126+
_14 = s // left_i # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
127+
check(assert_type(d // left_i, pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta)

0 commit comments

Comments
 (0)