diff --git a/pandas-stubs/core/arrays/datetimelike.pyi b/pandas-stubs/core/arrays/datetimelike.pyi index b5b5dba76..fc9548b0e 100644 --- a/pandas-stubs/core/arrays/datetimelike.pyi +++ b/pandas-stubs/core/arrays/datetimelike.pyi @@ -5,16 +5,21 @@ from pandas.core.arrays.base import ( ExtensionArray, ExtensionOpsMixin, ) +from typing_extensions import Self from pandas._libs import ( NaT as NaT, NaTType as NaTType, ) +from pandas._typing import TimeUnit class DatelikeOps: def strftime(self, date_format): ... class TimelikeOps: + @property + def unit(self) -> TimeUnit: ... + def as_unit(self, unit: TimeUnit) -> Self: ... def round(self, freq, ambiguous: str = ..., nonexistent: str = ...): ... def floor(self, freq, ambiguous: str = ..., nonexistent: str = ...): ... def ceil(self, freq, ambiguous: str = ..., nonexistent: str = ...): ... diff --git a/pandas-stubs/core/indexes/accessors.pyi b/pandas-stubs/core/indexes/accessors.pyi index a974fa6a9..437b8b195 100644 --- a/pandas-stubs/core/indexes/accessors.pyi +++ b/pandas-stubs/core/indexes/accessors.pyi @@ -36,6 +36,7 @@ from pandas._libs.tslibs import BaseOffset from pandas._libs.tslibs.offsets import DateOffset from pandas._typing import ( TimestampConvention, + TimeUnit, np_ndarray_bool, ) @@ -154,16 +155,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, @@ -173,7 +174,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]): | timedelta | Timedelta ) = ..., - ) -> _DTRoundingMethodReturnType: ... + ) -> _DTTimestampTimedeltaReturnType: ... def floor( self, freq: str | BaseOffset | None, @@ -183,7 +184,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]): | timedelta | Timedelta ) = ..., - ) -> _DTRoundingMethodReturnType: ... + ) -> _DTTimestampTimedeltaReturnType: ... def ceil( self, freq: str | BaseOffset | None, @@ -193,7 +194,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]): | timedelta | Timedelta ) = ..., - ) -> _DTRoundingMethodReturnType: ... + ) -> _DTTimestampTimedeltaReturnType: ... _DTNormalizeReturnType = TypeVar( "_DTNormalizeReturnType", TimestampSeries, DatetimeIndex @@ -202,9 +203,9 @@ _DTStrKindReturnType = TypeVar("_DTStrKindReturnType", Series[str], Index) _DTToPeriodReturnType = TypeVar("_DTToPeriodReturnType", PeriodSeries, PeriodIndex) class _DatetimeLikeNoTZMethods( - _DatetimeRoundingMethods[_DTRoundingMethodReturnType], + _DatetimeRoundingMethods[_DTTimestampTimedeltaReturnType], Generic[ - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTNormalizeReturnType, _DTStrKindReturnType, _DTToPeriodReturnType, @@ -238,7 +239,7 @@ class _DatetimeNoTZProperties( _DTFreqReturnType, ], _DatetimeLikeNoTZMethods[ - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTNormalizeReturnType, _DTStrKindReturnType, _DTToPeriodReturnType, @@ -246,7 +247,7 @@ class _DatetimeNoTZProperties( Generic[ _DTFieldOpsReturnType, _DTBoolOpsReturnType, - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTOtherOpsDateReturnType, _DTOtherOpsTimeReturnType, _DTFreqReturnType, @@ -261,7 +262,7 @@ class DatetimeProperties( _DatetimeNoTZProperties[ _DTFieldOpsReturnType, _DTBoolOpsReturnType, - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTOtherOpsDateReturnType, _DTOtherOpsTimeReturnType, _DTFreqReturnType, @@ -272,7 +273,7 @@ class DatetimeProperties( Generic[ _DTFieldOpsReturnType, _DTBoolOpsReturnType, - _DTRoundingMethodReturnType, + _DTTimestampTimedeltaReturnType, _DTOtherOpsDateReturnType, _DTOtherOpsTimeReturnType, _DTFreqReturnType, @@ -283,6 +284,9 @@ class DatetimeProperties( ): def to_pydatetime(self) -> np.ndarray: ... def isocalendar(self) -> DataFrame: ... + @property + def unit(self) -> TimeUnit: ... + def as_unit(self, unit: TimeUnit) -> _DTTimestampTimedeltaReturnType: ... _TDNoRoundingMethodReturnType = TypeVar( "_TDNoRoundingMethodReturnType", Series[int], Index @@ -309,7 +313,10 @@ class TimedeltaProperties( Properties, _TimedeltaPropertiesNoRounding[Series[int], Series[float]], _DatetimeRoundingMethods[TimedeltaSeries], -): ... +): + @property + def unit(self) -> TimeUnit: ... + def as_unit(self, unit: TimeUnit) -> TimedeltaSeries: ... _PeriodDTReturnTypes = TypeVar("_PeriodDTReturnTypes", TimestampSeries, DatetimeIndex) _PeriodIntReturnTypes = TypeVar("_PeriodIntReturnTypes", Series[int], Index[int]) diff --git a/pandas-stubs/core/indexes/datetimelike.pyi b/pandas-stubs/core/indexes/datetimelike.pyi index 5b138b744..fb27fae21 100644 --- a/pandas-stubs/core/indexes/datetimelike.pyi +++ b/pandas-stubs/core/indexes/datetimelike.pyi @@ -1,8 +1,12 @@ from pandas.core.indexes.extension import ExtensionIndex from pandas.core.indexes.timedeltas import TimedeltaIndex +from typing_extensions import Self from pandas._libs.tslibs import BaseOffset -from pandas._typing import S1 +from pandas._typing import ( + S1, + TimeUnit, +) class DatetimeIndexOpsMixin(ExtensionIndex[S1]): @property @@ -19,4 +23,7 @@ class DatetimeIndexOpsMixin(ExtensionIndex[S1]): self, other: DatetimeIndexOpsMixin ) -> TimedeltaIndex: ... -class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1]): ... +class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1]): + @property + def unit(self) -> TimeUnit: ... + def as_unit(self, unit: TimeUnit) -> Self: ... diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index eeb177b7c..36dd6d548 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -5,6 +5,7 @@ TYPE_CHECKING, Any, Optional, + cast, ) import numpy as np @@ -24,6 +25,8 @@ from pandas._typing import FulldatetimeDict else: FulldatetimeDict = Any +from pandas._typing import TimeUnit + from tests import ( TYPE_CHECKING_INVALID_USAGE, check, @@ -428,6 +431,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, TimeUnit), 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 +463,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, TimeUnit), 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, TimeUnit), 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, TimeUnit), 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: @@ -522,6 +559,11 @@ def test_datetimeindex_accessors() -> None: check(assert_type(i0.month_name(), pd.Index), pd.Index, str) check(assert_type(i0.day_name(), pd.Index), pd.Index, str) check(assert_type(i0.is_normalized, bool), bool) + check(assert_type(i0.unit, TimeUnit), str) + check(assert_type(i0.as_unit("s"), pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(i0.as_unit("ms"), pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(i0.as_unit("us"), pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(i0.as_unit("ns"), pd.DatetimeIndex), pd.DatetimeIndex) def test_timedeltaindex_accessors() -> None: @@ -542,6 +584,11 @@ def test_timedeltaindex_accessors() -> None: assert_type(i0.floor("D"), pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta ) check(assert_type(i0.ceil("D"), pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta) + check(assert_type(i0.unit, TimeUnit), str) + check(assert_type(i0.as_unit("s"), pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(i0.as_unit("ms"), pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(i0.as_unit("us"), pd.TimedeltaIndex), pd.TimedeltaIndex) + check(assert_type(i0.as_unit("ns"), pd.TimedeltaIndex), pd.TimedeltaIndex) def test_periodindex_accessors() -> None: