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

Author: loicdiridollou

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: (

tests/test_series.py

@@ -55,6 +55,10 @@
 )
 from tests.extension.decimal.array import DecimalDtype
 
+from pandas.core.arrays.datetimes import DatetimeArray
+
+from pandas.core.arrays.timedeltas import TimedeltaArray
+
 if TYPE_CHECKING:
     from pandas.core.series import (
         OffsetSeries,
@@ -577,7 +581,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 +3377,16 @@ 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)