Add unit property and as_unit method to DatetimeIndex, TimedeltaIndex and Series.dt (#865)

* Add unit property and as_unit method to DatetimeIndex and TimedeltaIndex

* Add unit property and as_unit method to TimelikeOps

* Add unit property and as_unit method to TimestampSeries and TimedeltaSeries

* Use TimeUnit type alias instead of str or Literal["s", "ms", "us", "ns"]

Author: skatsuta

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 = ...): ...

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,15 +239,15 @@ class _DatetimeNoTZProperties(
         _DTFreqReturnType,
     ],
     _DatetimeLikeNoTZMethods[
-        _DTRoundingMethodReturnType,
+        _DTTimestampTimedeltaReturnType,
         _DTNormalizeReturnType,
         _DTStrKindReturnType,
         _DTToPeriodReturnType,
     ],
     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])

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

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: