update mypy, pyright. handle slices of timestamps. Allow apply to return offsets.

pandas-stubs/_libs/tslibs/timestamps.pyi

@@ -11,6 +11,7 @@ from time import struct_time
 from typing import (
     ClassVar,
     Literal,
+    SupportsIndex,
     overload,
 )
 
@@ -48,7 +49,7 @@ _Nonexistent: TypeAlias = (
     Literal["raise", "NaT", "shift_backward", "shift_forward"] | Timedelta | timedelta
 )
 
-class Timestamp(datetime):
+class Timestamp(datetime, SupportsIndex):
     min: ClassVar[Timestamp]  # pyright: ignore[reportIncompatibleVariableOverride]
     max: ClassVar[Timestamp]  # pyright: ignore[reportIncompatibleVariableOverride]
 
@@ -309,3 +310,5 @@ class Timestamp(datetime):
     @property
     def unit(self) -> TimeUnit: ...
     def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...
+    # To support slicing
+    def __index__(self) -> int: ...

pandas-stubs/core/frame.pyi

@@ -9,6 +9,7 @@ from collections.abc import (
 )
 import datetime as dt
 from re import Pattern
+import sys
 from typing import (
     Any,
     ClassVar,
@@ -229,8 +230,32 @@ class _LocIndexerFrame(_LocIndexer):
         value: Scalar | NAType | NaTType | ArrayLike | Series | list | None,
     ) -> None: ...
 
-class DataFrame(NDFrame, OpsMixin):
-    __hash__: ClassVar[None]  # type: ignore[assignment]
+# With mypy 1.14.1 and python 3.12, the second overload needs a type-ignore statement
+if sys.version_info >= (3, 12):
+    class _GetItemHack:
+        @overload
+        def __getitem__(self, key: Scalar | tuple[Hashable, ...]) -> Series: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
+        @overload
+        def __getitem__(  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
+            self, key: Iterable[Hashable] | slice
+        ) -> DataFrame: ...
+        @overload
+        def __getitem__(self, key: Hashable) -> Series: ...
+
+else:
+    class _GetItemHack:
+        @overload
+        def __getitem__(self, key: Scalar | tuple[Hashable, ...]) -> Series: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
+        @overload
+        def __getitem__(  # pyright: ignore[reportOverlappingOverload]
+            self, key: Iterable[Hashable] | slice
+        ) -> DataFrame: ...
+        @overload
+        def __getitem__(self, key: Hashable) -> Series: ...
+
+class DataFrame(NDFrame, OpsMixin, _GetItemHack):
+
+    __hash__: ClassVar[None]  # type: ignore[assignment] # pyright: ignore[reportIncompatibleMethodOverride]
 
     @overload
     def __new__(
@@ -607,14 +632,6 @@ class DataFrame(NDFrame, OpsMixin):
     @property
     def T(self) -> DataFrame: ...
     def __getattr__(self, name: str) -> Series: ...
-    @overload
-    def __getitem__(self, key: Scalar | tuple[Hashable, ...]) -> Series: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
-    @overload
-    def __getitem__(  # pyright: ignore[reportOverlappingOverload]
-        self, key: Iterable[Hashable] | slice
-    ) -> DataFrame: ...
-    @overload
-    def __getitem__(self, key: Hashable) -> Series: ...
     def isetitem(
         self, loc: int | Sequence[int], value: Scalar | ArrayLike | list[Any]
     ) -> None: ...

pandas-stubs/core/series.pyi

@@ -1030,6 +1030,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
         **kwds,
     ) -> Series: ...
     @overload
+    def apply(
+        self,
+        func: Callable[..., BaseOffset],
+        convertDType: _bool = ...,
+        args: tuple = ...,
+        **kwds,
+    ) -> OffsetSeries: ...
+    @overload
     def apply(
         self,
         func: Callable[..., Series],

pyproject.toml

@@ -33,11 +33,11 @@ types-pytz = ">= 2022.1.1"
 numpy = ">= 1.23.5"
 
 [tool.poetry.group.dev.dependencies]
-mypy = "1.13.0"
+mypy = "1.14.1"
 pandas = "2.2.3"
 pyarrow = ">=10.0.1"
 pytest = ">=7.1.2"
-pyright = ">= 1.1.390"
+pyright = ">= 1.1.391"
 poethepoet = ">=0.16.5"
 loguru = ">=0.6.0"
 typing-extensions = ">=4.4.0"

tests/test_frame.py

@@ -2409,14 +2409,12 @@ def test_indexslice_getitem():
         .set_index(["x", "y"])
     )
     ind = pd.Index([2, 3])
-    # This next test is written this way to support both mypy 1.13 and newer
-    # versions of mypy and pyright that treat slice as a Generic due to
-    # a change in typeshed.
-    # Once pyright 1.1.390 and mypy 1.14 are released, the test can be
-    # reverted to the standard form.
-    # check(assert_type(pd.IndexSlice[ind, :], tuple["pd.Index[int]", slice]), tuple)
-    tmp = cast(tuple["pd.Index[int]", slice], pd.IndexSlice[ind, :])  # type: ignore[redundant-cast]
-    check(assert_type(tmp, tuple["pd.Index[int]", slice]), tuple)
+    check(
+        assert_type(
+            pd.IndexSlice[ind, :], tuple["pd.Index[int]", "slice[None, None, None]"]
+        ),
+        tuple,
+    )
     check(assert_type(df.loc[pd.IndexSlice[ind, :]], pd.DataFrame), pd.DataFrame)
     check(assert_type(df.loc[pd.IndexSlice[1:2]], pd.DataFrame), pd.DataFrame)
     check(

tests/test_series.py

@@ -3435,3 +3435,33 @@ def test_series_unique_timedelta() -> None:
     """Test type return of Series.unique on Series[timedeta64[ns]]."""
     sr = pd.Series([pd.Timedelta("1 days"), pd.Timedelta("3 days")])
     check(assert_type(sr.unique(), TimedeltaArray), TimedeltaArray)
+
+
+def test_slice_timestamp() -> None:
+    dti = pd.date_range("1/1/2025", "2/28/2025")
+
+    s = pd.Series([i for i in range(len(dti))], index=dti)
+
+    # For `s1`, see discussion in GH 397.  Needs mypy fix.
+    # s1 = s.loc["2025-01-15":"2025-01-20"]
+
+    # GH 397
+    check(
+        assert_type(
+            s.loc[pd.Timestamp("2025-01-15") : pd.Timestamp("2025-01-20")],
+            "pd.Series[int]",
+        ),
+        pd.Series,
+        np.integer,
+    )
+
+
+def test_apply_dateoffset() -> None:
+    # GH 454
+    months = [1, 2, 3]
+    s = pd.Series(months)
+    check(
+        assert_type(s.apply(lambda x: pd.DateOffset(months=x)), "OffsetSeries"),
+        pd.Series,
+        pd.DateOffset,
+    )