Skip to content

Commit 82f8994

Browse files
twoertweintwoertwein
authored
(r)true/floor div (#764)
* (r)true/floor div * address comments * and for DataFrame * add Sequence for TimedeltaIndex * list test for DataFrame * and 3.8 for pyright * avoid ZeroDivisionError * fix 32bit? * fix test on windows? * fix mypy on win? * use np.signedinteger
1 parent 7a8d310 commit 82f8994

File tree

9 files changed

+192
-63
lines changed

9 files changed

+192
-63
lines changed

‎pandas-stubs/core/arraylike.pyi

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,12 +29,12 @@ class OpsMixin:
2929
def __rsub__(self, other: Any) -> Self: ...
3030
def __mul__(self, other: Any) -> Self: ...
3131
def __rmul__(self, other: Any) -> Self: ...
32-
def __truediv__(self, other: Any) -> Self: ...
33-
def __rtruediv__(self, other: Any) -> Self: ...
34-
# __floordiv__ is handled by subclasses that specify only the valid values
32+
# Handled by subclasses that specify only the valid values
3533
# that can be passed
36-
# def __floordiv__(self, other: Any) -> Self: ...
37-
def __rfloordiv__(self, other: Any) -> Self: ...
34+
# def __truediv__(self, other: Any) -> Self: ...
35+
# def __rtruediv__(self, other: Any) -> Self: ...
36+
# def __floordiv__(self, other: Any) -> Self: ...
37+
# def __rfloordiv__(self, other: Any) -> Self: ...
3838
def __mod__(self, other: Any) -> Self: ...
3939
def __rmod__(self, other: Any) -> Self: ...
4040
def __divmod__(self, other: Any) -> tuple[Self, Self]: ...

‎pandas-stubs/core/frame.pyi

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2277,5 +2277,10 @@ class DataFrame(NDFrame, OpsMixin):
22772277
) -> DataFrame | Series: ...
22782278
# floordiv overload
22792279
def __floordiv__(
2280-
self, other: float | DataFrame | Series[int] | Series[float]
2280+
self, other: float | DataFrame | Series[int] | Series[float] | Sequence[float]
22812281
) -> Self: ...
2282+
def __rfloordiv__(
2283+
self, other: float | DataFrame | Series[int] | Series[float] | Sequence[float]
2284+
) -> Self: ...
2285+
def __truediv__(self, other: float | DataFrame | Series | Sequence) -> Self: ...
2286+
def __rtruediv__(self, other: float | DataFrame | Series | Sequence) -> Self: ...

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

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -430,17 +430,35 @@ class Index(IndexOpsMixin[S1]):
430430
def __floordiv__(
431431
self,
432432
other: float
433-
| Series[int]
434-
| Series[float]
433+
| IndexOpsMixin[int]
434+
| IndexOpsMixin[float]
435435
| Sequence[int]
436-
| Sequence[float]
437-
| Index[int]
438-
| Index[float],
436+
| Sequence[float],
437+
) -> Self: ...
438+
def __rfloordiv__(
439+
self,
440+
other: float
441+
| IndexOpsMixin[int]
442+
| IndexOpsMixin[float]
443+
| Sequence[int]
444+
| Sequence[float],
445+
) -> Self: ...
446+
def __truediv__(
447+
self,
448+
other: float
449+
| IndexOpsMixin[int]
450+
| IndexOpsMixin[float]
451+
| Sequence[int]
452+
| Sequence[float],
453+
) -> Self: ...
454+
def __rtruediv__(
455+
self,
456+
other: float
457+
| IndexOpsMixin[int]
458+
| IndexOpsMixin[float]
459+
| Sequence[int]
460+
| Sequence[float],
439461
) -> Self: ...
440-
@overload
441-
def __truediv__(self: Index[int] | Index[float], other: timedelta) -> Never: ...
442-
@overload
443-
def __truediv__(self, other: Any) -> Self: ...
444462

445463
def ensure_index_from_sequences(
446464
sequences: Sequence[Sequence[Dtype]], names: list[str] = ...

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ from typing import (
1111
import numpy as np
1212
from pandas import (
1313
DateOffset,
14+
Index,
1415
Period,
1516
)
1617
from pandas.core.indexes.accessors import TimedeltaIndexProperties
@@ -55,7 +56,20 @@ class TimedeltaIndex(DatetimeTimedeltaMixin[Timedelta], TimedeltaIndexProperties
5556
def __radd__(self, other: Timestamp | DatetimeIndex) -> DatetimeIndex: ... # type: ignore[override]
5657
def __sub__(self, other: Timedelta | Self) -> Self: ...
5758
def __mul__(self, other: num) -> Self: ...
58-
def __truediv__(self, other: num) -> Self: ...
59+
@overload # type: ignore[override]
60+
def __truediv__(self, other: num | Sequence[float]) -> Self: ...
61+
@overload
62+
def __truediv__(
63+
self, other: dt.timedelta | Sequence[dt.timedelta]
64+
) -> Index[float]: ...
65+
def __rtruediv__(self, other: dt.timedelta | Sequence[dt.timedelta]) -> Index[float]: ... # type: ignore[override]
66+
@overload # type: ignore[override]
67+
def __floordiv__(self, other: num | Sequence[float]) -> Self: ...
68+
@overload
69+
def __floordiv__(
70+
self, other: dt.timedelta | Sequence[dt.timedelta]
71+
) -> Index[int]: ...
72+
def __rfloordiv__(self, other: dt.timedelta | Sequence[dt.timedelta]) -> Index[int]: ... # type: ignore[override]
5973
def astype(self, dtype, copy: bool = ...): ...
6074
def get_value(self, series, key): ...
6175
def get_loc(self, key, tolerance=...): ...

‎pandas-stubs/core/series.pyi

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1502,9 +1502,6 @@ class Series(IndexOpsMixin[S1], NDFrame):
15021502
@overload
15031503
def __ror__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ... # type: ignore[misc]
15041504
def __rsub__(self, other: num | _ListLike | Series[S1]) -> Series: ...
1505-
@overload
1506-
def __rtruediv__(self, other: TimedeltaSeries) -> Series[float]: ...
1507-
@overload
15081505
def __rtruediv__(self, other: num | _ListLike | Series[S1]) -> Series: ...
15091506
# ignore needed for mypy as we want different results based on the arguments
15101507
@overload # type: ignore[override]
@@ -2030,7 +2027,44 @@ class TimedeltaSeries(Series[Timedelta]):
20302027
def __sub__( # type: ignore[override]
20312028
self, other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64
20322029
) -> TimedeltaSeries: ...
2033-
def __truediv__(self, other: Timedelta | TimedeltaSeries | np.timedelta64 | TimedeltaIndex) -> Series[float]: ... # type: ignore[override]
2030+
@overload # type: ignore[override]
2031+
def __truediv__(self, other: float | Sequence[float]) -> Self: ...
2032+
@overload
2033+
def __truediv__(
2034+
self,
2035+
other: timedelta
2036+
| TimedeltaSeries
2037+
| np.timedelta64
2038+
| TimedeltaIndex
2039+
| Sequence[timedelta],
2040+
) -> Series[float]: ...
2041+
def __rtruediv__( # type: ignore[override]
2042+
self,
2043+
other: timedelta
2044+
| TimedeltaSeries
2045+
| np.timedelta64
2046+
| TimedeltaIndex
2047+
| Sequence[timedelta],
2048+
) -> Series[float]: ...
2049+
@overload # type: ignore[override]
2050+
def __floordiv__(self, other: float | Sequence[float]) -> Self: ...
2051+
@overload
2052+
def __floordiv__(
2053+
self,
2054+
other: timedelta
2055+
| TimedeltaSeries
2056+
| np.timedelta64
2057+
| TimedeltaIndex
2058+
| Sequence[timedelta],
2059+
) -> Series[int]: ...
2060+
def __rfloordiv__( # type: ignore[override]
2061+
self,
2062+
other: timedelta
2063+
| TimedeltaSeries
2064+
| np.timedelta64
2065+
| TimedeltaIndex
2066+
| Sequence[timedelta],
2067+
) -> Series[int]: ...
20342068
@property
20352069
def dt(self) -> TimedeltaProperties: ... # type: ignore[override]
20362070
def mean( # type: ignore[override]

‎tests/test_frame.py

Lines changed: 49 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -745,59 +745,74 @@ def test_types_element_wise_arithmetic() -> None:
745745
df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
746746
df2 = pd.DataFrame(data={"col1": [10, 20], "col3": [3, 4]})
747747

748-
res_add1: pd.DataFrame = df + df2
749-
res_add2: pd.DataFrame = df.add(df2, fill_value=0)
748+
check(assert_type(df + df2, pd.DataFrame), pd.DataFrame)
749+
check(assert_type(df.add(df2, fill_value=0), pd.DataFrame), pd.DataFrame)
750750

751-
res_sub: pd.DataFrame = df - df2
752-
res_sub2: pd.DataFrame = df.sub(df2, fill_value=0)
751+
check(assert_type(df - df2, pd.DataFrame), pd.DataFrame)
752+
check(assert_type(df.sub(df2, fill_value=0), pd.DataFrame), pd.DataFrame)
753753

754-
res_mul: pd.DataFrame = df * df2
755-
res_mul2: pd.DataFrame = df.mul(df2, fill_value=0)
754+
check(assert_type(df * df2, pd.DataFrame), pd.DataFrame)
755+
check(assert_type(df.mul(df2, fill_value=0), pd.DataFrame), pd.DataFrame)
756756

757-
res_div: pd.DataFrame = df / df2
758-
res_div2: pd.DataFrame = df.div(df2, fill_value=0)
757+
check(assert_type(df / df2, pd.DataFrame), pd.DataFrame)
758+
check(assert_type(df.div(df2, fill_value=0), pd.DataFrame), pd.DataFrame)
759+
check(assert_type(df / [2, 2], pd.DataFrame), pd.DataFrame)
760+
check(assert_type(df.div([2, 2], fill_value=0), pd.DataFrame), pd.DataFrame)
759761

760-
res_floordiv: pd.DataFrame = df // df2
761-
res_floordiv2: pd.DataFrame = df.floordiv(df2, fill_value=0)
762+
check(assert_type(df // df2, pd.DataFrame), pd.DataFrame)
763+
check(assert_type(df.floordiv(df2, fill_value=0), pd.DataFrame), pd.DataFrame)
764+
check(assert_type(df // [2, 2], pd.DataFrame), pd.DataFrame)
765+
check(assert_type(df.floordiv([2, 2], fill_value=0), pd.DataFrame), pd.DataFrame)
762766

763-
res_mod: pd.DataFrame = df % df2
764-
res_mod2: pd.DataFrame = df.mod(df2, fill_value=0)
767+
check(assert_type(df % df2, pd.DataFrame), pd.DataFrame)
768+
check(assert_type(df.mod(df2, fill_value=0), pd.DataFrame), pd.DataFrame)
765769

766-
res_pow: pd.DataFrame = df2**df
767-
res_pow2: pd.DataFrame = df2.pow(df, fill_value=0)
770+
check(assert_type(df2**df, pd.DataFrame), pd.DataFrame)
771+
check(assert_type(df2.pow(df, fill_value=0), pd.DataFrame), pd.DataFrame)
768772

769773
# divmod operation was added in 1.2.0 https://pandas.pydata.org/docs/whatsnew/v1.2.0.html
770-
# noinspection PyTypeChecker
771-
res_divmod: tuple[pd.DataFrame, pd.DataFrame] = divmod(df, df2)
772-
res_divmod2: tuple[pd.DataFrame, pd.DataFrame] = df.__divmod__(df2)
773-
res_rdivmod: tuple[pd.DataFrame, pd.DataFrame] = df.__rdivmod__(df2)
774+
check(
775+
assert_type(divmod(df, df2), "tuple[pd.DataFrame, pd.DataFrame]"),
776+
tuple,
777+
pd.DataFrame,
778+
)
779+
check(
780+
assert_type(df.__divmod__(df2), "tuple[pd.DataFrame, pd.DataFrame]"),
781+
tuple,
782+
pd.DataFrame,
783+
)
784+
check(
785+
assert_type(df.__rdivmod__(df2), "tuple[pd.DataFrame, pd.DataFrame]"),
786+
tuple,
787+
pd.DataFrame,
788+
)
774789

775790

776791
def test_types_scalar_arithmetic() -> None:
777792
df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
778793

779-
res_add1: pd.DataFrame = df + 1
780-
res_add2: pd.DataFrame = df.add(1, fill_value=0)
794+
check(assert_type(df + 1, pd.DataFrame), pd.DataFrame)
795+
check(assert_type(df.add(1, fill_value=0), pd.DataFrame), pd.DataFrame)
781796

782-
res_sub: pd.DataFrame = df - 1
783-
res_sub2: pd.DataFrame = df.sub(1, fill_value=0)
797+
check(assert_type(df - 1, pd.DataFrame), pd.DataFrame)
798+
check(assert_type(df.sub(1, fill_value=0), pd.DataFrame), pd.DataFrame)
784799

785-
res_mul: pd.DataFrame = df * 2
786-
res_mul2: pd.DataFrame = df.mul(2, fill_value=0)
800+
check(assert_type(df * 2, pd.DataFrame), pd.DataFrame)
801+
check(assert_type(df.mul(2, fill_value=0), pd.DataFrame), pd.DataFrame)
787802

788-
res_div: pd.DataFrame = df / 2
789-
res_div2: pd.DataFrame = df.div(2, fill_value=0)
803+
check(assert_type(df / 2, pd.DataFrame), pd.DataFrame)
804+
check(assert_type(df.div(2, fill_value=0), pd.DataFrame), pd.DataFrame)
790805

791-
res_floordiv: pd.DataFrame = df // 2
792-
res_floordiv2: pd.DataFrame = df.floordiv(2, fill_value=0)
806+
check(assert_type(df // 2, pd.DataFrame), pd.DataFrame)
807+
check(assert_type(df.floordiv(2, fill_value=0), pd.DataFrame), pd.DataFrame)
793808

794-
res_mod: pd.DataFrame = df % 2
795-
res_mod2: pd.DataFrame = df.mod(2, fill_value=0)
809+
check(assert_type(df % 2, pd.DataFrame), pd.DataFrame)
810+
check(assert_type(df.mod(2, fill_value=0), pd.DataFrame), pd.DataFrame)
796811

797-
res_pow: pd.DataFrame = df**2
798-
res_pow1: pd.DataFrame = df**0
799-
res_pow2: pd.DataFrame = df**0.213
800-
res_pow3: pd.DataFrame = df.pow(0.5)
812+
check(assert_type(df**2, pd.DataFrame), pd.DataFrame)
813+
check(assert_type(df**0, pd.DataFrame), pd.DataFrame)
814+
check(assert_type(df**0.213, pd.DataFrame), pd.DataFrame)
815+
check(assert_type(df.pow(0.5), pd.DataFrame), pd.DataFrame)
801816

802817

803818
def test_types_melt() -> None:

‎tests/test_indexes.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1022,3 +1022,28 @@ def test_new() -> None:
10221022
pd.IntervalIndex,
10231023
pd.Interval,
10241024
)
1025+
1026+
1027+
def test_timedelta_div() -> None:
1028+
index = pd.Index([pd.Timedelta(days=1)], dtype="timedelta64[s]")
1029+
delta = dt.timedelta(1)
1030+
1031+
check(assert_type(index / delta, "pd.Index[float]"), pd.Index, float)
1032+
check(assert_type(index / [delta], "pd.Index[float]"), pd.Index, float)
1033+
check(assert_type(index / 1, pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta)
1034+
check(assert_type(index / [1], pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta)
1035+
check(assert_type(index // delta, "pd.Index[int]"), pd.Index, np.longlong)
1036+
check(assert_type(index // [delta], "pd.Index[int]"), pd.Index, int)
1037+
check(assert_type(index // 1, pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta)
1038+
check(assert_type(index // [1], pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta)
1039+
1040+
check(assert_type(delta / index, "pd.Index[float]"), pd.Index, float)
1041+
check(assert_type([delta] / index, "pd.Index[float]"), pd.Index, float)
1042+
check(assert_type(delta // index, "pd.Index[int]"), pd.Index, np.longlong)
1043+
check(assert_type([delta] // index, "pd.Index[int]"), pd.Index, np.signedinteger)
1044+
1045+
if TYPE_CHECKING_INVALID_USAGE:
1046+
1 / index # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]
1047+
[1] / index # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]
1048+
1 // index # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]
1049+
[1] // index # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]

‎tests/test_scalars.py

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
import pandas as pd
1616
import pytz
1717
from typing_extensions import (
18-
Never,
1918
TypeAlias,
2019
assert_type,
2120
)
@@ -780,14 +779,8 @@ def test_timedelta_mul_div() -> None:
780779
md_ndarray_float / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]
781780
mp_series_int / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]
782781
md_series_float / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]
783-
assert_type(
784-
md_int64_index / td, # pyright: ignore[reportGeneralTypeIssues]
785-
Never,
786-
)
787-
assert_type(
788-
md_float_index / td, # pyright: ignore[reportGeneralTypeIssues]
789-
Never,
790-
)
782+
md_int64_index / td, # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]
783+
md_float_index / td, # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]
791784

792785

793786
def test_timedelta_mod_abs_unary() -> None:

‎tests/test_series.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2686,3 +2686,28 @@ def double(x):
26862686
# Test cases with None and pd.NA as other
26872687
check(assert_type(s.mask(s > 3, None), pd.Series), pd.Series, np.float64)
26882688
check(assert_type(s.mask(s > 3, pd.NA), pd.Series), pd.Series, np.float64)
2689+
2690+
2691+
def test_timedelta_div() -> None:
2692+
series = pd.Series([pd.Timedelta(days=1)])
2693+
delta = datetime.timedelta(1)
2694+
2695+
check(assert_type(series / delta, "pd.Series[float]"), pd.Series, float)
2696+
check(assert_type(series / [delta], "pd.Series[float]"), pd.Series, float)
2697+
check(assert_type(series / 1, "TimedeltaSeries"), pd.Series, pd.Timedelta)
2698+
check(assert_type(series / [1], "TimedeltaSeries"), pd.Series, pd.Timedelta)
2699+
check(assert_type(series // delta, "pd.Series[int]"), pd.Series, np.longlong)
2700+
check(assert_type(series // [delta], "pd.Series[int]"), pd.Series, int)
2701+
check(assert_type(series // 1, "TimedeltaSeries"), pd.Series, pd.Timedelta)
2702+
check(assert_type(series // [1], "TimedeltaSeries"), pd.Series, pd.Timedelta)
2703+
2704+
check(assert_type(delta / series, "pd.Series[float]"), pd.Series, float)
2705+
check(assert_type([delta] / series, "pd.Series[float]"), pd.Series, float)
2706+
check(assert_type(delta // series, "pd.Series[int]"), pd.Series, np.longlong)
2707+
check(assert_type([delta] // series, "pd.Series[int]"), pd.Series, np.signedinteger)
2708+
2709+
if TYPE_CHECKING_INVALID_USAGE:
2710+
1 / series # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]
2711+
[1] / series # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]
2712+
1 // series # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]
2713+
[1] // series # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]

0 commit comments

Comments
 (0)