Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ from pandas._typing import (
FilePath,
FillnaOptions,
FormattersType,
Frequency,
GroupByObjectNonScalar,
HashableT,
HashableT1,
Expand Down Expand Up @@ -855,7 +856,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
def shift(
self,
periods: int = ...,
freq=...,
freq: Frequency | dt.timedelta | None = ...,
axis: Axis = ...,
fill_value: Hashable | None = ...,
) -> Self: ...
Expand Down
3 changes: 2 additions & 1 deletion pandas-stubs/core/series.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,7 @@ from pandas._typing import (
FilePath,
FillnaOptions,
FloatDtypeArg,
Frequency,
GroupByObjectNonScalar,
HashableT1,
IgnoreRaise,
Expand Down Expand Up @@ -1219,7 +1220,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
def shift(
self,
periods: int = ...,
freq=...,
freq: Frequency | timedelta | None = ...,
axis: AxisIndex = ...,
fill_value: object | None = ...,
) -> Series[S1]: ...
Expand Down
11 changes: 7 additions & 4 deletions tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -516,10 +516,13 @@ def test_types_sort_values_with_key() -> None:


def test_types_shift() -> None:
df = pd.DataFrame(data={"col1": [1, 1], "col2": [3, 4]})
df.shift()
df.shift(1)
df.shift(-1)
df = pd.DataFrame(
data={"col1": [1, 1], "col2": [3, 4]}, index=pd.date_range("2020", periods=2)
)
check(assert_type(df.shift(), pd.DataFrame), pd.DataFrame)
check(assert_type(df.shift(1), pd.DataFrame), pd.DataFrame)
check(assert_type(df.shift(-1), pd.DataFrame), pd.DataFrame)
check(assert_type(df.shift(freq="1D"), pd.DataFrame), pd.DataFrame)


def test_types_rank() -> None:
Expand Down
18 changes: 14 additions & 4 deletions tests/test_series.py
Original file line number Diff line number Diff line change
Expand Up @@ -416,10 +416,20 @@ def test_types_sort_values_with_key() -> None:


def test_types_shift() -> None:
s = pd.Series([1, 2, 3])
s.shift()
s.shift(axis=0, periods=1)
s.shift(-1, fill_value=0)
s = pd.Series([1, 2, 3], index=pd.date_range("2020", periods=3))
# Return type "Series[int]"" not quite correct
# - https://github.com/pandas-dev/pandas-stubs/issues/1111
# - https://github.com/pandas-dev/pandas-stubs/issues/1110
check(assert_type(s.shift(), "pd.Series[int]"), pd.Series, np.floating)
check(
assert_type(s.shift(axis=0, periods=1), "pd.Series[int]"),
pd.Series,
np.floating,
)
check(
assert_type(s.shift(-1, fill_value=0), "pd.Series[int]"), pd.Series, np.integer
)
check(assert_type(s.shift(freq="1D"), "pd.Series[int]"), pd.Series, np.integer)


def test_types_rank() -> None:
Expand Down