diff --git a/pandas-stubs/core/indexes/interval.pyi b/pandas-stubs/core/indexes/interval.pyi index c1474174b..aeef461fe 100644 --- a/pandas-stubs/core/indexes/interval.pyi +++ b/pandas-stubs/core/indexes/interval.pyi @@ -371,7 +371,7 @@ def interval_range( start: _TimestampLike, end: _TimestampLike = ..., periods: int | None = ..., - freq: str | BaseOffset | None = ..., + freq: str | BaseOffset | pd.Timedelta | dt.timedelta | None = ..., name: Hashable = ..., closed: IntervalClosedType = ..., ) -> IntervalIndex[Interval[pd.Timestamp]]: ... @@ -381,7 +381,7 @@ def interval_range( start: None = ..., end: _TimestampLike, periods: int | None = ..., - freq: str | BaseOffset | None = ..., + freq: str | BaseOffset | pd.Timedelta | dt.timedelta | None = ..., name: Hashable = ..., closed: IntervalClosedType = ..., ) -> IntervalIndex[Interval[pd.Timestamp]]: ... @@ -391,7 +391,7 @@ def interval_range( *, end: None = ..., periods: int | None = ..., - freq: str | BaseOffset | None = ..., + freq: str | BaseOffset | pd.Timedelta | dt.timedelta | None = ..., name: Hashable = ..., closed: IntervalClosedType = ..., ) -> IntervalIndex[Interval[pd.Timestamp]]: ... @@ -400,7 +400,7 @@ def interval_range( start: _TimedeltaLike, end: _TimedeltaLike = ..., periods: int | None = ..., - freq: str | BaseOffset | None = ..., + freq: str | BaseOffset | pd.Timedelta | dt.timedelta | None = ..., name: Hashable = ..., closed: IntervalClosedType = ..., ) -> IntervalIndex[Interval[pd.Timedelta]]: ... @@ -410,7 +410,7 @@ def interval_range( start: None = ..., end: _TimedeltaLike, periods: int | None = ..., - freq: str | BaseOffset | None = ..., + freq: str | BaseOffset | pd.Timedelta | dt.timedelta | None = ..., name: Hashable = ..., closed: IntervalClosedType = ..., ) -> IntervalIndex[Interval[pd.Timedelta]]: ... @@ -420,7 +420,7 @@ def interval_range( *, end: None = ..., periods: int | None = ..., - freq: str | BaseOffset | None = ..., + freq: str | BaseOffset | pd.Timedelta | dt.timedelta | None = ..., name: Hashable = ..., closed: IntervalClosedType = ..., ) -> IntervalIndex[Interval[pd.Timedelta]]: ... diff --git a/pandas-stubs/core/indexes/timedeltas.pyi b/pandas-stubs/core/indexes/timedeltas.pyi index 545898190..9a58bdc0d 100644 --- a/pandas-stubs/core/indexes/timedeltas.pyi +++ b/pandas-stubs/core/indexes/timedeltas.pyi @@ -72,7 +72,7 @@ def timedelta_range( start: TimedeltaConvertibleTypes = ..., end: TimedeltaConvertibleTypes = ..., periods: int | None = ..., - freq: str | DateOffset | None = ..., + freq: str | DateOffset | Timedelta | dt.timedelta | None = ..., name: Hashable | None = ..., closed: Literal["left", "right"] | None = ..., ) -> TimedeltaIndex: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index c53c288bb..41c1e9f27 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -179,7 +179,13 @@ class _LocIndexerSeries(_LocIndexer, Generic[S1]): @overload def __getitem__( self, - idx: MaskType | Index | Sequence[float] | list[str] | slice | _IndexSliceTuple, + idx: MaskType + | Index + | Sequence[float] + | list[str] + | slice + | _IndexSliceTuple + | Callable, # _IndexSliceTuple is when having a tuple that includes a slice. Could just # be s.loc[1, :], or s.loc[pd.IndexSlice[1, :]] ) -> Series[S1]: ... @@ -775,7 +781,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def apply( self, - func: Callable[..., Scalar | Sequence | set | Mapping], + func: Callable[..., Scalar | Sequence | set | Mapping | None], convertDType: _bool = ..., args: tuple = ..., **kwds, diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 2c337a35d..fa4557600 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -291,6 +291,30 @@ def test_interval_range(): pd.IntervalIndex, pd.Interval, ) + check( + assert_type( + pd.interval_range( + pd.Timestamp(2000, 1, 1), + pd.Timestamp(2010, 1, 1), + freq=pd.Timedelta(days=30), + ), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", + ), + pd.IntervalIndex, + pd.Interval, + ) + check( + assert_type( + pd.interval_range( + pd.Timestamp(2000, 1, 1), + pd.Timestamp(2010, 1, 1), + freq=dt.timedelta(days=30), + ), + "pd.IntervalIndex[pd.Interval[pd.Timestamp]]", + ), + pd.IntervalIndex, + pd.Interval, + ) check( assert_type( pd.interval_range(pd.Timestamp(2000, 1, 1), dt.datetime(2010, 1, 1), 5), @@ -308,6 +332,26 @@ def test_interval_range(): pd.IntervalIndex, pd.Interval, ) + check( + assert_type( + pd.interval_range( + pd.Timedelta("1D"), pd.Timedelta("10D"), freq=pd.Timedelta("2D") + ), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", + ), + pd.IntervalIndex, + pd.Interval, + ) + check( + assert_type( + pd.interval_range( + pd.Timedelta("1D"), pd.Timedelta("10D"), freq=dt.timedelta(days=2) + ), + "pd.IntervalIndex[pd.Interval[pd.Timedelta]]", + ), + pd.IntervalIndex, + pd.Interval, + ) check( assert_type( pd.interval_range(end=pd.Timedelta("10D"), periods=10, freq="D"), diff --git a/tests/test_series.py b/tests/test_series.py index d37d0726a..f76c2662d 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -1820,3 +1820,15 @@ def test_convert_dtypes_dtype_backend() -> None: s = pd.Series([1, 2, 3, 4]) s1 = s.convert_dtypes(dtype_backend="numpy_nullable") check(assert_type(s1, pd.Series), pd.Series) + + +def test_apply_returns_none() -> None: + # GH 557 + s = pd.Series([1, 2, 3]) + check(assert_type(s.apply(lambda x: None), pd.Series), pd.Series) + + +def test_loc_callable() -> None: + # GH 586 + s = pd.Series([1, 2]) + check(assert_type(s.loc[lambda x: x > 1], pd.Series), pd.Series) diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index fe5c80677..e1d6e329c 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -1071,6 +1071,20 @@ def test_timedelta_range() -> None: ), pd.TimedeltaIndex, ) + check( + assert_type( + pd.timedelta_range("1 day", "10 days", freq=pd.Timedelta("2 days")), + pd.TimedeltaIndex, + ), + pd.TimedeltaIndex, + ) + check( + assert_type( + pd.timedelta_range("1 day", "10 days", freq=dt.timedelta(days=2)), + pd.TimedeltaIndex, + ), + pd.TimedeltaIndex, + ) def test_dateoffset_freqstr() -> None: