Skip to content

Commit

Permalink
GH1039 Clear tests and add typehint for Series.unique with datetime/t…
Browse files Browse the repository at this point in the history
…imedelta (#1040)

* GH1039 Clear tests and add typehint for Series.unique with datetime/timedelta

* Formatting
  • Loading branch information
loicdiridollou authored Nov 21, 2024
1 parent 779aab6 commit e610b76
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 1 deletion.
4 changes: 4 additions & 0 deletions pandas-stubs/core/series.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,10 @@ from pandas.core.api import (
Int32Dtype as Int32Dtype,
Int64Dtype as Int64Dtype,
)
from pandas.core.arrays import TimedeltaArray
from pandas.core.arrays.base import ExtensionArray
from pandas.core.arrays.categorical import CategoricalAccessor
from pandas.core.arrays.datetimes import DatetimeArray
from pandas.core.arrays.interval import IntervalArray
from pandas.core.base import IndexOpsMixin
from pandas.core.frame import DataFrame
Expand Down Expand Up @@ -2113,6 +2115,7 @@ class TimestampSeries(Series[Timestamp]):
) -> TimestampSeries: ...
def __mul__(self, other: float | Series[int] | Series[float] | Sequence[float]) -> TimestampSeries: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
def __truediv__(self, other: float | Series[int] | Series[float] | Sequence[float]) -> TimestampSeries: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
def unique(self) -> DatetimeArray: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
def mean( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
self,
axis: AxisIndex | None = ...,
Expand Down Expand Up @@ -2156,6 +2159,7 @@ class TimedeltaSeries(Series[Timedelta]):
def __mul__( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
self, other: num | Sequence[num] | Series[int] | Series[float]
) -> TimedeltaSeries: ...
def unique(self) -> TimedeltaArray: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
def __sub__( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
self,
other: (
Expand Down
16 changes: 15 additions & 1 deletion tests/test_series.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@
ExtensionArray,
ExtensionDtype,
)
from pandas.core.arrays.datetimes import DatetimeArray
from pandas.core.arrays.timedeltas import TimedeltaArray
from pandas.core.window import ExponentialMovingWindow
import pytest
from typing_extensions import (
Expand Down Expand Up @@ -577,7 +579,7 @@ def test_types_value_counts() -> None:

def test_types_unique() -> None:
s = pd.Series([-10, 2, 2, 3, 10, 10])
s.unique()
check(assert_type(s.unique(), np.ndarray), np.ndarray)


def test_types_apply() -> None:
Expand Down Expand Up @@ -3373,3 +3375,15 @@ def test_case_when() -> None:
]
)
check(assert_type(r, pd.Series), pd.Series)


def test_series_unique_timestamp() -> None:
"""Test type return of Series.unique on Series[datetime64[ns]]."""
sr = pd.Series(pd.bdate_range("2023-10-10", "2023-10-15"))
check(assert_type(sr.unique(), DatetimeArray), DatetimeArray)


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)

0 comments on commit e610b76

Please sign in to comment.