make sub of timedelta work (#759)

* make sub of timedelta work

* fix issues due to pyright 1.1.320

* fix for older python with list subscript

* change order of test so ruff does not remove import from typeshed

* get ruff to ignore a line

* add comments about why testing is special for mypy vs pyright

Author: Dr-Irv

pandas-stubs/core/series.pyi

@@ -1533,6 +1533,16 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     @overload
     def __rxor__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ...  # type: ignore[misc]
     @overload
+    def __sub__(
+        self: Series[Timestamp],
+        other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64,
+    ) -> TimestampSeries: ...
+    @overload
+    def __sub__(
+        self: Series[Timedelta],
+        other: Timedelta | TimedeltaSeries | TimedeltaIndex | np.timedelta64,
+    ) -> TimedeltaSeries: ...
+    @overload
     def __sub__(
         self, other: Timestamp | datetime | TimestampSeries
     ) -> TimedeltaSeries: ...

pyproject.toml

@@ -197,6 +197,7 @@ reportPrivateUsage = false
 # enable optional checks
 reportMissingModuleSource = true
 useLibraryCodeForTypes = false
+defineConstant = { MYPY_CHECKING = false }
 
 [tool.codespell]
 ignore-words-list = "indext, mose, sav, ser"

tests/test_indexes.py

@@ -24,6 +24,17 @@
 )
 
 if TYPE_CHECKING:
+    MYPY_CHECKING: bool = True
+    # See test_sorted_and_list() where mypy and pyright do different
+    # inference on sorted(pd.Index)
+    if MYPY_CHECKING:
+        from typing import Any
+
+        from typing_extensions import TypeAlias
+
+        SupportsRichComparison: TypeAlias = Any
+    else:
+        from _typeshed import SupportsRichComparison  # noqa: F401
     from pandas.core.indexes.base import (
         _ComplexIndexType,
         _FloatIndexType,
@@ -719,10 +730,12 @@ def test_interval_index_tuples():
 def test_sorted_and_list() -> None:
     # GH 497
     i1 = pd.Index([3, 2, 1])
+    # mypy infers sorted(i1) as list[Any], while pyright infers sorted(i1) as
+    # list[SupportsRichComparison]
     check(
         assert_type(
             sorted(i1),
-            list,
+            "list[SupportsRichComparison]",
         ),
         list,
     )

tests/test_pandas.py

@@ -1272,7 +1272,7 @@ def test_merge_ordered() -> None:
         )
         pd.merge_ordered(  # type: ignore[call-overload]
             ls,
-            rf,  # pyright: ignore[reportGeneralTypeIssues]
+            rf,
             left_on="left",
             right_on="b",
             left_by="left",  # pyright: ignore[reportGeneralTypeIssues]

tests/test_timefuncs.py

@@ -1130,7 +1130,14 @@ 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
-    # https://github.com/pandas-dev/pandas/issues/54059 needs to be fixed
+    check(assert_type(s3, "TimedeltaSeries"), pd.Series, pd.Timedelta)
+    td1 = pd.Timedelta(1, "D")
+    check(assert_type(s2 - td1, "TimestampSeries"), pd.Series, pd.Timestamp)
+    # GH 758
+    s4 = s1.astype(object)
+    check(assert_type(s4 - td1, "TimestampSeries"), pd.Series, pd.Timestamp)
+
+    # https://github.com/pandas-dev/pandas/issues/54059 says this is invalid
     if PD_LTE_20:
         td = np.timedelta64(1, "M")
         check(assert_type((s1 - td), "TimestampSeries"), pd.Series, pd.Timestamp)