diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/pandas-stubs/_libs/tslibs/timestamps.pyi index 504b801cc..d99661b17 100644 --- a/pandas-stubs/_libs/tslibs/timestamps.pyi +++ b/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: ... diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 10f127b19..1bb3b4e54 100644 --- a/pandas-stubs/core/frame.pyi +++ b/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: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index b91d3842b..6d104138d 100644 --- a/pandas-stubs/core/series.pyi +++ b/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], diff --git a/pyproject.toml b/pyproject.toml index 11a3dd93b..c233df4c9 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/test_frame.py b/tests/test_frame.py index 5ce029cfd..48230cf3b 100644 --- a/tests/test_frame.py +++ b/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( diff --git a/tests/test_series.py b/tests/test_series.py index 87055fb74..f710b931f 100644 --- a/tests/test_series.py +++ b/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, + )