fix: type asserts

attempt.py

@@ -0,0 +1,48 @@
+from typing import assert_type, reveal_type
+import pandas as pd
+from pandas.core.series import TimedeltaSeries  # noqa: F401
+import numpy as np
+import datetime as dt
+
+from tests import check
+
+
+df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [5, 4, 3, 2, 1]})
+s1 = df.min(axis=1)
+s2 = df.max(axis=1)
+sa = s1 + s2
+ss = s1 - s2
+sm = s1 * s2
+sd = s1 / s2
+check(assert_type(sa, pd.Series), pd.Series)
+reveal_type(s1.__sub__(s2))
+reveal_type(s2.__rsub__(s1))
+check(assert_type(ss, pd.Series), pd.Series)
+check(assert_type(sm, pd.Series), pd.Series)
+check(assert_type(sd, pd.Series), pd.Series)
+
+ts1 = pd.to_datetime(pd.Series(["2022-03-05", "2022-03-06"]))
+assert isinstance(ts1.iloc[0], pd.Timestamp)
+td1 = pd.to_timedelta([2, 3], "seconds")
+ts2 = pd.to_datetime(pd.Series(["2022-03-08", "2022-03-10"]))
+r1 = ts1 - ts2
+check(assert_type(r1, "TimedeltaSeries"), pd.Series, pd.Timedelta)
+r2 = r1 / td1
+check(assert_type(r2, "pd.Series[float]"), pd.Series, float)
+r3 = r1 - td1
+check(assert_type(r3, "TimedeltaSeries"), pd.Series, pd.Timedelta)
+r4 = pd.Timedelta(5, "days") / r1
+check(assert_type(r4, "pd.Series[float]"), pd.Series, float)
+sb = pd.Series([1, 2]) == pd.Series([1, 3])
+check(assert_type(sb, "pd.Series[bool]"), pd.Series, np.bool_)
+r5 = sb * r1
+check(assert_type(r5, "TimedeltaSeries"), pd.Series, pd.Timedelta)
+r6 = r1 * 4
+check(assert_type(r6, "TimedeltaSeries"), pd.Series, pd.Timedelta)
+
+tsp1 = pd.Timestamp("2022-03-05")
+dt1 = dt.datetime(2022, 9, 1, 12, 5, 30)
+r7 = ts1 - tsp1
+check(assert_type(r7, "TimedeltaSeries"), pd.Series, pd.Timedelta)
+r8 = ts1 - dt1
+check(assert_type(r8, "TimedeltaSeries"), pd.Series, pd.Timedelta)

pandas-stubs/_libs/tslibs/timestamps.pyi

@@ -246,9 +246,9 @@ class Timestamp(datetime, SupportsIndex):
     @overload
     def __sub__(self, other: TimedeltaSeries) -> Series[Timestamp]: ...
     @overload
-    def __sub__(self, other: Series[Never]) -> Series: ...  # type: ignore[overload-overlap]
+    def __sub__(self, other: Series[Never]) -> Series: ...
     @overload
-    def __sub__(self, other: Series[Timestamp]) -> TimedeltaSeries: ...
+    def __sub__(self, other: Series[Timestamp]) -> Series[Timedelta]: ...
     @overload
     def __sub__(
         self, other: np_ndarray[ShapeT, np.timedelta64]

pandas-stubs/core/series.pyi

@@ -1624,9 +1624,9 @@ class Series(IndexOpsMixin[S1], NDFrame):
     # just failed to generate these so I couldn't match
     # them up.
     @overload
-    def __add__(self: Series[Never], other: Scalar | _ListLike | Series) -> Series: ...
+    def __add__(self: Series[Never], other: Scalar | _ListLike | Series) -> Series: ...  # type: ignore[overload-overlap]
     @overload
-    def __add__(self, other: Series[Never]) -> Series: ...
+    def __add__(self, other: Series[Never]) -> Series: ...  # type: ignore[overload-overlap]
     @overload
     def __add__(
         self: Series[bool],
@@ -1709,6 +1709,15 @@ class Series(IndexOpsMixin[S1], NDFrame):
         other: datetime | np.datetime64 | np_ndarray_dt | Series[Timestamp],
     ) -> Series[Timestamp]: ...
     @overload
+    def __add__(
+        self: Series[Timedelta],
+        other: timedelta | np.timedelta64 | np_ndarray_td | TimedeltaSeries,
+    ) -> TimedeltaSeries: ...
+    @overload
+    def __add__(
+        self: Series[Timedelta], other: Series[Timedelta]
+    ) -> Series[Timedelta]: ...
+    @overload
     def __add__(self: Series[Timedelta], other: Period) -> PeriodSeries: ...
     @overload
     def add(
@@ -2235,16 +2244,38 @@ class Series(IndexOpsMixin[S1], NDFrame):
         other: timedelta | np.timedelta64 | np_ndarray_td | TimedeltaSeries,
     ) -> TimedeltaSeries: ...
     @overload
+    def __mul__(self: Series[bool], other: Series[Timedelta]) -> Series[Timedelta]: ...  # type: ignore[overload-overlap]
+    @overload
     def __mul__(
         self: Series[int],
         other: timedelta | np.timedelta64 | np_ndarray_td | TimedeltaSeries,
     ) -> TimedeltaSeries: ...
     @overload
+    def __mul__(self: Series[int], other: Series[Timedelta]) -> Series[Timedelta]: ...
+    @overload
     def __mul__(
         self: Series[float],
         other: timedelta | np.timedelta64 | np_ndarray_td | TimedeltaSeries,
     ) -> TimedeltaSeries: ...
     @overload
+    def __mul__(self: Series[float], other: Series[Timedelta]) -> Series[Timedelta]: ...
+    @overload
+    def __mul__(
+        self: Series[Timedelta],
+        other: (
+            float
+            | Sequence[float]
+            | np_ndarray_bool
+            | np_ndarray_anyint
+            | np_ndarray_float
+        ),
+    ) -> TimedeltaSeries: ...
+    @overload
+    def __mul__(
+        self: Series[Timedelta],
+        other: Series[bool] | Series[int] | Series[float],
+    ) -> Series[Timedelta]: ...
+    @overload
     def mul(
         self: Series[Never],
         other: complex | _ListLike | Series,
@@ -2436,19 +2467,49 @@ class Series(IndexOpsMixin[S1], NDFrame):
         self: Series[_T_COMPLEX], other: np_ndarray_complex
     ) -> Series[complex]: ...
     @overload
-    def __rmul__(
+    def __rmul__(  # type: ignore[misc]
         self: Series[bool],
-        other: timedelta | np.timedelta64 | np_ndarray_td | TimedeltaSeries,
+        other: (
+            timedelta
+            | np.timedelta64
+            | np_ndarray_td
+            | Series[Timedelta]
+            | TimedeltaSeries
+        ),
     ) -> TimedeltaSeries: ...
     @overload
-    def __rmul__(
+    def __rmul__(  # type: ignore[misc]
         self: Series[int],
-        other: timedelta | np.timedelta64 | np_ndarray_td | TimedeltaSeries,
+        other: (
+            timedelta
+            | np.timedelta64
+            | np_ndarray_td
+            | Series[Timedelta]
+            | TimedeltaSeries
+        ),
     ) -> TimedeltaSeries: ...
     @overload
-    def __rmul__(
+    def __rmul__(  # type: ignore[misc]
         self: Series[float],
-        other: timedelta | np.timedelta64 | np_ndarray_td | TimedeltaSeries,
+        other: (
+            timedelta
+            | np.timedelta64
+            | np_ndarray_td
+            | Series[Timedelta]
+            | TimedeltaSeries
+        ),
+    ) -> TimedeltaSeries: ...
+    @overload
+    def __rmul__(
+        self: Series[Timedelta],
+        other: (
+            float
+            | Sequence[float]
+            | np_ndarray_bool
+            | np_ndarray_anyint
+            | np_ndarray_float
+            | Series[_T_INT]
+        ),
     ) -> TimedeltaSeries: ...
     @overload
     def rmul(
@@ -2627,16 +2688,17 @@ class Series(IndexOpsMixin[S1], NDFrame):
     @overload
     def __rxor__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ...
     @overload
-    def __sub__(self, other: Series[Never]) -> Series: ...  # type: ignore[overload-overlap]
+    def __sub__(self: Series[Never], other: Series[Never]) -> Series: ...  # type: ignore[overload-overlap]
     @overload
     def __sub__(
-        self: Series[Never],
-        other: datetime | np.datetime64 | np_ndarray_dt | Series[Timestamp],
+        self: Series[Never], other: datetime | np.datetime64 | np_ndarray_dt
     ) -> TimedeltaSeries: ...
     @overload
-    def __sub__(  # type: ignore[overload-overlap]
-        self: Series[Never], other: complex | _ListLike | Series
-    ) -> Series: ...
+    def __sub__(self: Series[Never], other: Series[Timestamp]) -> Series[Timedelta]: ...
+    @overload
+    def __sub__(self: Series[Never], other: complex | _ListLike | Series) -> Series: ...
+    @overload
+    def __sub__(self, other: Series[Never]) -> Series: ...  # type: ignore[overload-overlap]
     @overload
     def __sub__(
         self: Series[bool],
@@ -2702,18 +2764,20 @@ class Series(IndexOpsMixin[S1], NDFrame):
     ) -> Series[complex]: ...
     @overload
     def __sub__(
-        self: Series[Timestamp],
-        other: datetime | np.datetime64 | np_ndarray_dt | Series[Timestamp],
+        self: Series[Timestamp], other: datetime | np.datetime64 | np_ndarray_dt
     ) -> TimedeltaSeries: ...
     @overload
+    def __sub__(
+        self: Series[Timestamp], other: Series[Timestamp]
+    ) -> Series[Timedelta]: ...
+    @overload
     def __sub__(
         self: Series[Timestamp],
         other: (
             timedelta
             | np.timedelta64
             | np_ndarray_td
             | TimedeltaIndex
-            | Series[Timedelta]
             | TimedeltaSeries
             | BaseOffset
         ),
@@ -2726,7 +2790,6 @@ class Series(IndexOpsMixin[S1], NDFrame):
             | np.timedelta64
             | np_ndarray_td
             | TimedeltaIndex
-            | Series[Timedelta]
             | TimedeltaSeries
         ),
     ) -> TimedeltaSeries: ...

tests/series/arithmetic/test_sub.py

@@ -180,7 +180,7 @@ def test_sub_pd_datetime() -> None:
     a = pd.Series([s + pd.Timedelta(minutes=m) for m in range(3)])
 
     check(assert_type(left_ts - s, "TimedeltaSeries"), pd.Series, pd.Timedelta)
-    check(assert_type(left_ts - a, "TimedeltaSeries"), pd.Series, pd.Timedelta)
+    check(assert_type(left_ts - a, "pd.Series[pd.Timedelta]"), pd.Series, pd.Timedelta)
 
     check(assert_type(s - left_ts, pd.Series), pd.Series, pd.Timedelta)
     check(assert_type(a - left_ts, pd.Series), pd.Series, pd.Timedelta)

tests/series/test_series.py

@@ -1598,6 +1598,7 @@ def test_series_min_max_sub_axis() -> None:
     ss = s1 - s2
     sm = s1 * s2
     sd = s1 / s2
+    s1.__sub__(s2)
     check(assert_type(sa, pd.Series), pd.Series)
     check(assert_type(ss, pd.Series), pd.Series)
     check(assert_type(sm, pd.Series), pd.Series)

tests/test_timefuncs.py

@@ -191,7 +191,7 @@ def test_timestamp_timedelta_series_arithmetic() -> None:
     td1 = pd.to_timedelta([2, 3], "seconds")
     ts2 = pd.to_datetime(pd.Series(["2022-03-08", "2022-03-10"]))
     r1 = ts1 - ts2
-    check(assert_type(r1, "TimedeltaSeries"), pd.Series, pd.Timedelta)
+    check(assert_type(r1, "pd.Series[pd.Timedelta]"), pd.Series, pd.Timedelta)
     r2 = r1 / td1
     check(assert_type(r2, "pd.Series[float]"), pd.Series, float)
     r3 = r1 - td1
@@ -201,7 +201,7 @@ def test_timestamp_timedelta_series_arithmetic() -> None:
     sb = pd.Series([1, 2]) == pd.Series([1, 3])
     check(assert_type(sb, "pd.Series[bool]"), pd.Series, np.bool_)
     r5 = sb * r1
-    check(assert_type(r5, "TimedeltaSeries"), pd.Series, pd.Timedelta)
+    check(assert_type(r5, "pd.Series[pd.Timedelta]"), pd.Series, pd.Timedelta)
     r6 = r1 * 4
     check(assert_type(r6, "TimedeltaSeries"), pd.Series, pd.Timedelta)
 
@@ -1655,7 +1655,7 @@ def test_timedelta64_and_arithmatic_operator() -> None:
     s1 = pd.Series(data=pd.date_range("1/1/2020", "2/1/2020"))
     s2 = pd.Series(data=pd.date_range("1/1/2021", "2/1/2021"))
     s3 = s2 - s1
-    check(assert_type(s3, "TimedeltaSeries"), pd.Series, pd.Timedelta)
+    check(assert_type(s3, "pd.Series[pd.Timedelta]"), pd.Series, pd.Timedelta)
     td1 = pd.Timedelta(1, "D")
     check(assert_type(s2 - td1, "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
     # GH 758
@@ -1808,7 +1808,7 @@ def test_timestamp_sub_series() -> None:
     ts1 = pd.to_datetime(pd.Series(["2022-03-05", "2022-03-06"]))
     one_ts = ts1.iloc[0]
     check(assert_type(ts1.iloc[0], pd.Timestamp), pd.Timestamp)
-    check(assert_type(one_ts - ts1, "TimedeltaSeries"), pd.Series, pd.Timedelta)
+    check(assert_type(one_ts - ts1, "pd.Series[pd.Timedelta]"), pd.Series, pd.Timedelta)
 
 
 def test_creating_date_range() -> None: