From cd35a2f427b2fb227ec2db83128414b8029ea79b Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Mon, 12 Feb 2024 15:48:04 +0900 Subject: [PATCH] Add unit property and as_unit method to TimestampSeries and TimedeltaSeries --- pandas-stubs/core/indexes/accessors.pyi | 36 +++++++++++++++---------- tests/test_timefuncs.py | 35 ++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 14 deletions(-) diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index 1accd842..e6ad88bb 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/pandas-stubs/core/indexes/accessors.pyi @@ -154,16 +154,16 @@ class _DatetimeLikeOps( # type of the series, we don't know which kind of series was ...ed # in to the dt accessor -_DTRoundingMethodReturnType = TypeVar( - "_DTRoundingMethodReturnType", +_DTTimestampTimedeltaReturnType = TypeVar( + "_DTTimestampTimedeltaReturnType", Series, - TimedeltaSeries, TimestampSeries, + TimedeltaSeries, DatetimeIndex, TimedeltaIndex, ) -class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]): +class _DatetimeRoundingMethods(Generic[_DTTimestampTimedeltaReturnType]): def round( self, freq: str | BaseOffset | None, @@ -171,7 +171,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]): nonexistent: Literal["shift_forward", "shift_backward", "NaT", "raise"] | timedelta | Timedelta = ..., - ) -> _DTRoundingMethodReturnType: ... + ) -> _DTTimestampTimedeltaReturnType: ... def floor( self, freq: str | BaseOffset | None, @@ -179,7 +179,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]): nonexistent: Literal["shift_forward", "shift_backward", "NaT", "raise"] | timedelta | Timedelta = ..., - ) -> _DTRoundingMethodReturnType: ... + ) -> _DTTimestampTimedeltaReturnType: ... def ceil( self, freq: str | BaseOffset | None, @@ -187,7 +187,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]): nonexistent: Literal["shift_forward", "shift_backward", "NaT", "raise"] | timedelta | Timedelta = ..., - ) -> _DTRoundingMethodReturnType: ... + ) -> _DTTimestampTimedeltaReturnType: ... _DTNormalizeReturnType = TypeVar( "_DTNormalizeReturnType", TimestampSeries, DatetimeIndex @@ -196,9 +196,9 @@ _DTStrKindReturnType = TypeVar("_DTStrKindReturnType", Series[str], Index) _DTToPeriodReturnType = TypeVar("_DTToPeriodReturnType", PeriodSeries, PeriodIndex) class _DatetimeLikeNoTZMethods( - _DatetimeRoundingMethods[_DTRoundingMethodReturnType], + _DatetimeRoundingMethods[_DTTimestampTimedeltaReturnType], Generic[ - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTNormalizeReturnType, _DTStrKindReturnType, _DTToPeriodReturnType, @@ -230,7 +230,7 @@ class _DatetimeNoTZProperties( _DTFreqReturnType, ], _DatetimeLikeNoTZMethods[ - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTNormalizeReturnType, _DTStrKindReturnType, _DTToPeriodReturnType, @@ -238,7 +238,7 @@ class _DatetimeNoTZProperties( Generic[ _DTFieldOpsReturnType, _DTBoolOpsReturnType, - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTOtherOpsDateReturnType, _DTOtherOpsTimeReturnType, _DTFreqReturnType, @@ -253,7 +253,7 @@ class DatetimeProperties( _DatetimeNoTZProperties[ _DTFieldOpsReturnType, _DTBoolOpsReturnType, - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTOtherOpsDateReturnType, _DTOtherOpsTimeReturnType, _DTFreqReturnType, @@ -264,7 +264,7 @@ class DatetimeProperties( Generic[ _DTFieldOpsReturnType, _DTBoolOpsReturnType, - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTOtherOpsDateReturnType, _DTOtherOpsTimeReturnType, _DTFreqReturnType, @@ -275,6 +275,11 @@ class DatetimeProperties( ): def to_pydatetime(self) -> np.ndarray: ... def isocalendar(self) -> DataFrame: ... + @property + def unit(self) -> str: ... + def as_unit( + self, unit: Literal["s", "ms", "us", "ns"] + ) -> _DTTimestampTimedeltaReturnType: ... _TDNoRoundingMethodReturnType = TypeVar( "_TDNoRoundingMethodReturnType", Series[int], Index @@ -301,7 +306,10 @@ class TimedeltaProperties( Properties, _TimedeltaPropertiesNoRounding[Series[int], Series[float]], _DatetimeRoundingMethods[TimedeltaSeries], -): ... +): + @property + def unit(self) -> str: ... + def as_unit(self, unit: Literal["s", "ms", "us", "ns"]) -> TimedeltaSeries: ... _PeriodDTReturnTypes = TypeVar("_PeriodDTReturnTypes", TimestampSeries, DatetimeIndex) _PeriodIntReturnTypes = TypeVar("_PeriodIntReturnTypes", Series[int], Index[int]) diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index a50edaa4..ca6b8d20 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -5,6 +5,7 @@ TYPE_CHECKING, Any, Optional, + cast, ) import numpy as np @@ -428,6 +429,11 @@ def test_series_dt_accessors() -> None: ) check(assert_type(s0.dt.month_name(), "pd.Series[str]"), pd.Series, str) check(assert_type(s0.dt.day_name(), "pd.Series[str]"), pd.Series, str) + check(assert_type(s0.dt.unit, str), str) + check(assert_type(s0.dt.as_unit("s"), "TimestampSeries"), pd.Series, pd.Timestamp) + check(assert_type(s0.dt.as_unit("ms"), "TimestampSeries"), pd.Series, pd.Timestamp) + check(assert_type(s0.dt.as_unit("us"), "TimestampSeries"), pd.Series, pd.Timestamp) + check(assert_type(s0.dt.as_unit("ns"), "TimestampSeries"), pd.Series, pd.Timestamp) i1 = pd.period_range(start="2022-06-01", periods=10) @@ -455,6 +461,35 @@ def test_series_dt_accessors() -> None: check(assert_type(s2.dt.components, pd.DataFrame), pd.DataFrame) check(assert_type(s2.dt.to_pytimedelta(), np.ndarray), np.ndarray) check(assert_type(s2.dt.total_seconds(), "pd.Series[float]"), pd.Series, float) + check(assert_type(s2.dt.unit, str), str) + check(assert_type(s2.dt.as_unit("s"), "TimedeltaSeries"), pd.Series, pd.Timedelta) + check(assert_type(s2.dt.as_unit("ms"), "TimedeltaSeries"), pd.Series, pd.Timedelta) + check(assert_type(s2.dt.as_unit("us"), "TimedeltaSeries"), pd.Series, pd.Timedelta) + check(assert_type(s2.dt.as_unit("ns"), "TimedeltaSeries"), pd.Series, pd.Timedelta) + + # Checks for general Series other than TimestampSeries and TimedeltaSeries + + s4 = cast( + "pd.Series[pd.Timestamp]", + pd.Series([pd.Timestamp("2024-01-01"), pd.Timestamp("2024-01-02")]), + ) + + check(assert_type(s4.dt.unit, str), str) + check(assert_type(s4.dt.as_unit("s"), pd.Series), pd.Series, pd.Timestamp) + check(assert_type(s4.dt.as_unit("ms"), pd.Series), pd.Series, pd.Timestamp) + check(assert_type(s4.dt.as_unit("us"), pd.Series), pd.Series, pd.Timestamp) + check(assert_type(s4.dt.as_unit("ns"), pd.Series), pd.Series, pd.Timestamp) + + s5 = cast( + "pd.Series[pd.Timedelta]", + pd.Series([pd.Timedelta("1 day"), pd.Timedelta("2 days")]), + ) + + check(assert_type(s5.dt.unit, str), str) + check(assert_type(s5.dt.as_unit("s"), pd.Series), pd.Series, pd.Timedelta) + check(assert_type(s5.dt.as_unit("ms"), pd.Series), pd.Series, pd.Timedelta) + check(assert_type(s5.dt.as_unit("us"), pd.Series), pd.Series, pd.Timedelta) + check(assert_type(s5.dt.as_unit("ns"), pd.Series), pd.Series, pd.Timedelta) def test_datetimeindex_accessors() -> None: