diff --git a/pandas-stubs/core/arraylike.pyi b/pandas-stubs/core/arraylike.pyi index 9c8d95f43..08d47a123 100644 --- a/pandas-stubs/core/arraylike.pyi +++ b/pandas-stubs/core/arraylike.pyi @@ -29,12 +29,12 @@ class OpsMixin: def __rsub__(self, other: Any) -> Self: ... def __mul__(self, other: Any) -> Self: ... def __rmul__(self, other: Any) -> Self: ... - def __truediv__(self, other: Any) -> Self: ... - def __rtruediv__(self, other: Any) -> Self: ... - # __floordiv__ is handled by subclasses that specify only the valid values + # Handled by subclasses that specify only the valid values # that can be passed - # def __floordiv__(self, other: Any) -> Self: ... - def __rfloordiv__(self, other: Any) -> Self: ... + # def __truediv__(self, other: Any) -> Self: ... + # def __rtruediv__(self, other: Any) -> Self: ... + # def __floordiv__(self, other: Any) -> Self: ... + # def __rfloordiv__(self, other: Any) -> Self: ... def __mod__(self, other: Any) -> Self: ... def __rmod__(self, other: Any) -> Self: ... def __divmod__(self, other: Any) -> tuple[Self, Self]: ... diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 71e00c245..152d97f18 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -2277,5 +2277,10 @@ class DataFrame(NDFrame, OpsMixin): ) -> DataFrame | Series: ... # floordiv overload def __floordiv__( - self, other: float | DataFrame | Series[int] | Series[float] + self, other: float | DataFrame | Series[int] | Series[float] | Sequence[float] ) -> Self: ... + def __rfloordiv__( + self, other: float | DataFrame | Series[int] | Series[float] | Sequence[float] + ) -> Self: ... + def __truediv__(self, other: float | DataFrame | Series | Sequence) -> Self: ... + def __rtruediv__(self, other: float | DataFrame | Series | Sequence) -> Self: ... diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 95f1277e7..2dc604047 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -430,17 +430,35 @@ class Index(IndexOpsMixin[S1]): def __floordiv__( self, other: float - | Series[int] - | Series[float] + | IndexOpsMixin[int] + | IndexOpsMixin[float] | Sequence[int] - | Sequence[float] - | Index[int] - | Index[float], + | Sequence[float], + ) -> Self: ... + def __rfloordiv__( + self, + other: float + | IndexOpsMixin[int] + | IndexOpsMixin[float] + | Sequence[int] + | Sequence[float], + ) -> Self: ... + def __truediv__( + self, + other: float + | IndexOpsMixin[int] + | IndexOpsMixin[float] + | Sequence[int] + | Sequence[float], + ) -> Self: ... + def __rtruediv__( + self, + other: float + | IndexOpsMixin[int] + | IndexOpsMixin[float] + | Sequence[int] + | Sequence[float], ) -> Self: ... - @overload - def __truediv__(self: Index[int] | Index[float], other: timedelta) -> Never: ... - @overload - def __truediv__(self, other: Any) -> Self: ... def ensure_index_from_sequences( sequences: Sequence[Sequence[Dtype]], names: list[str] = ... diff --git a/pandas-stubs/core/indexes/timedeltas.pyi b/pandas-stubs/core/indexes/timedeltas.pyi index bcf480ee3..4cdf732c9 100644 --- a/pandas-stubs/core/indexes/timedeltas.pyi +++ b/pandas-stubs/core/indexes/timedeltas.pyi @@ -11,6 +11,7 @@ from typing import ( import numpy as np from pandas import ( DateOffset, + Index, Period, ) from pandas.core.indexes.accessors import TimedeltaIndexProperties @@ -55,7 +56,20 @@ class TimedeltaIndex(DatetimeTimedeltaMixin[Timedelta], TimedeltaIndexProperties def __radd__(self, other: Timestamp | DatetimeIndex) -> DatetimeIndex: ... # type: ignore[override] def __sub__(self, other: Timedelta | Self) -> Self: ... def __mul__(self, other: num) -> Self: ... - def __truediv__(self, other: num) -> Self: ... + @overload # type: ignore[override] + def __truediv__(self, other: num | Sequence[float]) -> Self: ... + @overload + def __truediv__( + self, other: dt.timedelta | Sequence[dt.timedelta] + ) -> Index[float]: ... + def __rtruediv__(self, other: dt.timedelta | Sequence[dt.timedelta]) -> Index[float]: ... # type: ignore[override] + @overload # type: ignore[override] + def __floordiv__(self, other: num | Sequence[float]) -> Self: ... + @overload + def __floordiv__( + self, other: dt.timedelta | Sequence[dt.timedelta] + ) -> Index[int]: ... + def __rfloordiv__(self, other: dt.timedelta | Sequence[dt.timedelta]) -> Index[int]: ... # type: ignore[override] def astype(self, dtype, copy: bool = ...): ... def get_value(self, series, key): ... def get_loc(self, key, tolerance=...): ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 839c2a19a..c1055c48e 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1502,9 +1502,6 @@ class Series(IndexOpsMixin[S1], NDFrame): @overload def __ror__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ... # type: ignore[misc] def __rsub__(self, other: num | _ListLike | Series[S1]) -> Series: ... - @overload - def __rtruediv__(self, other: TimedeltaSeries) -> Series[float]: ... - @overload def __rtruediv__(self, other: num | _ListLike | Series[S1]) -> Series: ... # ignore needed for mypy as we want different results based on the arguments @overload # type: ignore[override] @@ -2030,7 +2027,44 @@ class TimedeltaSeries(Series[Timedelta]): def __sub__( # type: ignore[override] self, other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64 ) -> TimedeltaSeries: ... - def __truediv__(self, other: Timedelta | TimedeltaSeries | np.timedelta64 | TimedeltaIndex) -> Series[float]: ... # type: ignore[override] + @overload # type: ignore[override] + def __truediv__(self, other: float | Sequence[float]) -> Self: ... + @overload + def __truediv__( + self, + other: timedelta + | TimedeltaSeries + | np.timedelta64 + | TimedeltaIndex + | Sequence[timedelta], + ) -> Series[float]: ... + def __rtruediv__( # type: ignore[override] + self, + other: timedelta + | TimedeltaSeries + | np.timedelta64 + | TimedeltaIndex + | Sequence[timedelta], + ) -> Series[float]: ... + @overload # type: ignore[override] + def __floordiv__(self, other: float | Sequence[float]) -> Self: ... + @overload + def __floordiv__( + self, + other: timedelta + | TimedeltaSeries + | np.timedelta64 + | TimedeltaIndex + | Sequence[timedelta], + ) -> Series[int]: ... + def __rfloordiv__( # type: ignore[override] + self, + other: timedelta + | TimedeltaSeries + | np.timedelta64 + | TimedeltaIndex + | Sequence[timedelta], + ) -> Series[int]: ... @property def dt(self) -> TimedeltaProperties: ... # type: ignore[override] def mean( # type: ignore[override] diff --git a/tests/test_frame.py b/tests/test_frame.py index e577be021..b71880032 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -745,59 +745,74 @@ def test_types_element_wise_arithmetic() -> None: df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]}) df2 = pd.DataFrame(data={"col1": [10, 20], "col3": [3, 4]}) - res_add1: pd.DataFrame = df + df2 - res_add2: pd.DataFrame = df.add(df2, fill_value=0) + check(assert_type(df + df2, pd.DataFrame), pd.DataFrame) + check(assert_type(df.add(df2, fill_value=0), pd.DataFrame), pd.DataFrame) - res_sub: pd.DataFrame = df - df2 - res_sub2: pd.DataFrame = df.sub(df2, fill_value=0) + check(assert_type(df - df2, pd.DataFrame), pd.DataFrame) + check(assert_type(df.sub(df2, fill_value=0), pd.DataFrame), pd.DataFrame) - res_mul: pd.DataFrame = df * df2 - res_mul2: pd.DataFrame = df.mul(df2, fill_value=0) + check(assert_type(df * df2, pd.DataFrame), pd.DataFrame) + check(assert_type(df.mul(df2, fill_value=0), pd.DataFrame), pd.DataFrame) - res_div: pd.DataFrame = df / df2 - res_div2: pd.DataFrame = df.div(df2, fill_value=0) + check(assert_type(df / df2, pd.DataFrame), pd.DataFrame) + check(assert_type(df.div(df2, fill_value=0), pd.DataFrame), pd.DataFrame) + check(assert_type(df / [2, 2], pd.DataFrame), pd.DataFrame) + check(assert_type(df.div([2, 2], fill_value=0), pd.DataFrame), pd.DataFrame) - res_floordiv: pd.DataFrame = df // df2 - res_floordiv2: pd.DataFrame = df.floordiv(df2, fill_value=0) + check(assert_type(df // df2, pd.DataFrame), pd.DataFrame) + check(assert_type(df.floordiv(df2, fill_value=0), pd.DataFrame), pd.DataFrame) + check(assert_type(df // [2, 2], pd.DataFrame), pd.DataFrame) + check(assert_type(df.floordiv([2, 2], fill_value=0), pd.DataFrame), pd.DataFrame) - res_mod: pd.DataFrame = df % df2 - res_mod2: pd.DataFrame = df.mod(df2, fill_value=0) + check(assert_type(df % df2, pd.DataFrame), pd.DataFrame) + check(assert_type(df.mod(df2, fill_value=0), pd.DataFrame), pd.DataFrame) - res_pow: pd.DataFrame = df2**df - res_pow2: pd.DataFrame = df2.pow(df, fill_value=0) + check(assert_type(df2**df, pd.DataFrame), pd.DataFrame) + check(assert_type(df2.pow(df, fill_value=0), pd.DataFrame), pd.DataFrame) # divmod operation was added in 1.2.0 https://pandas.pydata.org/docs/whatsnew/v1.2.0.html - # noinspection PyTypeChecker - res_divmod: tuple[pd.DataFrame, pd.DataFrame] = divmod(df, df2) - res_divmod2: tuple[pd.DataFrame, pd.DataFrame] = df.__divmod__(df2) - res_rdivmod: tuple[pd.DataFrame, pd.DataFrame] = df.__rdivmod__(df2) + check( + assert_type(divmod(df, df2), "tuple[pd.DataFrame, pd.DataFrame]"), + tuple, + pd.DataFrame, + ) + check( + assert_type(df.__divmod__(df2), "tuple[pd.DataFrame, pd.DataFrame]"), + tuple, + pd.DataFrame, + ) + check( + assert_type(df.__rdivmod__(df2), "tuple[pd.DataFrame, pd.DataFrame]"), + tuple, + pd.DataFrame, + ) def test_types_scalar_arithmetic() -> None: df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]}) - res_add1: pd.DataFrame = df + 1 - res_add2: pd.DataFrame = df.add(1, fill_value=0) + check(assert_type(df + 1, pd.DataFrame), pd.DataFrame) + check(assert_type(df.add(1, fill_value=0), pd.DataFrame), pd.DataFrame) - res_sub: pd.DataFrame = df - 1 - res_sub2: pd.DataFrame = df.sub(1, fill_value=0) + check(assert_type(df - 1, pd.DataFrame), pd.DataFrame) + check(assert_type(df.sub(1, fill_value=0), pd.DataFrame), pd.DataFrame) - res_mul: pd.DataFrame = df * 2 - res_mul2: pd.DataFrame = df.mul(2, fill_value=0) + check(assert_type(df * 2, pd.DataFrame), pd.DataFrame) + check(assert_type(df.mul(2, fill_value=0), pd.DataFrame), pd.DataFrame) - res_div: pd.DataFrame = df / 2 - res_div2: pd.DataFrame = df.div(2, fill_value=0) + check(assert_type(df / 2, pd.DataFrame), pd.DataFrame) + check(assert_type(df.div(2, fill_value=0), pd.DataFrame), pd.DataFrame) - res_floordiv: pd.DataFrame = df // 2 - res_floordiv2: pd.DataFrame = df.floordiv(2, fill_value=0) + check(assert_type(df // 2, pd.DataFrame), pd.DataFrame) + check(assert_type(df.floordiv(2, fill_value=0), pd.DataFrame), pd.DataFrame) - res_mod: pd.DataFrame = df % 2 - res_mod2: pd.DataFrame = df.mod(2, fill_value=0) + check(assert_type(df % 2, pd.DataFrame), pd.DataFrame) + check(assert_type(df.mod(2, fill_value=0), pd.DataFrame), pd.DataFrame) - res_pow: pd.DataFrame = df**2 - res_pow1: pd.DataFrame = df**0 - res_pow2: pd.DataFrame = df**0.213 - res_pow3: pd.DataFrame = df.pow(0.5) + check(assert_type(df**2, pd.DataFrame), pd.DataFrame) + check(assert_type(df**0, pd.DataFrame), pd.DataFrame) + check(assert_type(df**0.213, pd.DataFrame), pd.DataFrame) + check(assert_type(df.pow(0.5), pd.DataFrame), pd.DataFrame) def test_types_melt() -> None: diff --git a/tests/test_indexes.py b/tests/test_indexes.py index be437f5e0..0115b152d 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -1022,3 +1022,28 @@ def test_new() -> None: pd.IntervalIndex, pd.Interval, ) + + +def test_timedelta_div() -> None: + index = pd.Index([pd.Timedelta(days=1)], dtype="timedelta64[s]") + delta = dt.timedelta(1) + + check(assert_type(index / delta, "pd.Index[float]"), pd.Index, float) + check(assert_type(index / [delta], "pd.Index[float]"), pd.Index, float) + check(assert_type(index / 1, pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta) + check(assert_type(index / [1], pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta) + check(assert_type(index // delta, "pd.Index[int]"), pd.Index, np.longlong) + check(assert_type(index // [delta], "pd.Index[int]"), pd.Index, int) + check(assert_type(index // 1, pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta) + check(assert_type(index // [1], pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta) + + check(assert_type(delta / index, "pd.Index[float]"), pd.Index, float) + check(assert_type([delta] / index, "pd.Index[float]"), pd.Index, float) + check(assert_type(delta // index, "pd.Index[int]"), pd.Index, np.longlong) + check(assert_type([delta] // index, "pd.Index[int]"), pd.Index, np.signedinteger) + + if TYPE_CHECKING_INVALID_USAGE: + 1 / index # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + [1] / index # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + 1 // index # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + [1] // index # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 652814fec..9692676b2 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -15,7 +15,6 @@ import pandas as pd import pytz from typing_extensions import ( - Never, TypeAlias, assert_type, ) @@ -780,14 +779,8 @@ def test_timedelta_mul_div() -> None: md_ndarray_float / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] mp_series_int / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] md_series_float / td # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] - assert_type( - md_int64_index / td, # pyright: ignore[reportGeneralTypeIssues] - Never, - ) - assert_type( - md_float_index / td, # pyright: ignore[reportGeneralTypeIssues] - Never, - ) + md_int64_index / td, # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + md_float_index / td, # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] def test_timedelta_mod_abs_unary() -> None: diff --git a/tests/test_series.py b/tests/test_series.py index 88c9f4854..a9eced036 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -2686,3 +2686,28 @@ def double(x): # Test cases with None and pd.NA as other check(assert_type(s.mask(s > 3, None), pd.Series), pd.Series, np.float64) check(assert_type(s.mask(s > 3, pd.NA), pd.Series), pd.Series, np.float64) + + +def test_timedelta_div() -> None: + series = pd.Series([pd.Timedelta(days=1)]) + delta = datetime.timedelta(1) + + check(assert_type(series / delta, "pd.Series[float]"), pd.Series, float) + check(assert_type(series / [delta], "pd.Series[float]"), pd.Series, float) + check(assert_type(series / 1, "TimedeltaSeries"), pd.Series, pd.Timedelta) + check(assert_type(series / [1], "TimedeltaSeries"), pd.Series, pd.Timedelta) + check(assert_type(series // delta, "pd.Series[int]"), pd.Series, np.longlong) + check(assert_type(series // [delta], "pd.Series[int]"), pd.Series, int) + check(assert_type(series // 1, "TimedeltaSeries"), pd.Series, pd.Timedelta) + check(assert_type(series // [1], "TimedeltaSeries"), pd.Series, pd.Timedelta) + + check(assert_type(delta / series, "pd.Series[float]"), pd.Series, float) + check(assert_type([delta] / series, "pd.Series[float]"), pd.Series, float) + check(assert_type(delta // series, "pd.Series[int]"), pd.Series, np.longlong) + check(assert_type([delta] // series, "pd.Series[int]"), pd.Series, np.signedinteger) + + if TYPE_CHECKING_INVALID_USAGE: + 1 / series # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + [1] / series # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + 1 // series # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues] + [1] // series # type: ignore[operator] # pyright: ignore[reportGeneralTypeIssues]