From 09bb8734ed880791b0a4887b8fde8486fefb2050 Mon Sep 17 00:00:00 2001 From: Loic Diridollou Date: Wed, 20 Nov 2024 18:26:47 -0500 Subject: [PATCH 1/2] GH1039 Clear tests and add typehint for Series.unique with datetime/timedelta --- pandas-stubs/core/series.pyi | 4 ++++ tests/test_series.py | 19 ++++++++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) 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..54fd72731 100644 --- a/tests/test_series.py +++ b/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) From d7dbf67c4248877305855a7c13b9b64e07f8e262 Mon Sep 17 00:00:00 2001 From: Loic Diridollou Date: Wed, 20 Nov 2024 18:28:41 -0500 Subject: [PATCH 2/2] Formatting --- tests/test_series.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/test_series.py b/tests/test_series.py index 54fd72731..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 ( @@ -55,10 +57,6 @@ ) 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, @@ -3381,11 +3379,10 @@ def test_case_when() -> None: 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')) + 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")])