diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 457db8de3..692c0cdde 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -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 @@ -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 = ..., @@ -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: ( diff --git a/tests/test_series.py b/tests/test_series.py index 8e9b60bde..147844a24 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -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 ( @@ -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: @@ -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)