From e610b7646a5c1f2c5c18c90b4934377e12652d58 Mon Sep 17 00:00:00 2001
From: Loic Diridollou <loic.diridollou@gmail.com>
Date: Thu, 21 Nov 2024 09:46:06 -0500
Subject: [PATCH] GH1039 Clear tests and add typehint for Series.unique with
 datetime/timedelta (#1040)

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

* Formatting
---
 pandas-stubs/core/series.pyi |  4 ++++
 tests/test_series.py         | 16 +++++++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi
index 457db8de..692c0cdd 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 8e9b60bd..147844a2 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)