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

Author: skatsuta

pandas-stubs/core/arrays/datetimelike.pyi

@@ -1,5 +1,4 @@
 from collections.abc import Sequence
-from typing import Literal
 
 import numpy as np
 from pandas.core.arrays.base import (
@@ -12,14 +11,15 @@ 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) -> str: ...
-    def as_unit(self, unit: Literal["s", "ms", "us", "ns"]) -> Self: ...
+    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,
 )
 
@@ -284,10 +285,8 @@ 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: ...
+    def unit(self) -> TimeUnit: ...
+    def as_unit(self, unit: TimeUnit) -> _DTTimestampTimedeltaReturnType: ...
 
 _TDNoRoundingMethodReturnType = TypeVar(
     "_TDNoRoundingMethodReturnType", Series[int], Index
@@ -316,8 +315,8 @@ class TimedeltaProperties(
     _DatetimeRoundingMethods[TimedeltaSeries],
 ):
     @property
-    def unit(self) -> str: ...
-    def as_unit(self, unit: Literal["s", "ms", "us", "ns"]) -> TimedeltaSeries: ...
+    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,11 +1,12 @@
-from typing import Literal
-
 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
@@ -24,5 +25,5 @@ class DatetimeIndexOpsMixin(ExtensionIndex[S1]):
 
 class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1]):
     @property
-    def unit(self) -> str: ...
-    def as_unit(self, unit: Literal["s", "ms", "us", "ns"]) -> Self: ...
+    def unit(self) -> TimeUnit: ...
+    def as_unit(self, unit: TimeUnit) -> Self: ...

tests/test_timefuncs.py

@@ -25,6 +25,8 @@
     from pandas._typing import FulldatetimeDict
 else:
     FulldatetimeDict = Any
+from pandas._typing import TimeUnit
+
 from tests import (
     TYPE_CHECKING_INVALID_USAGE,
     check,
@@ -429,7 +431,7 @@ 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.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)
@@ -461,7 +463,7 @@ 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.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)
@@ -474,7 +476,7 @@ def test_series_dt_accessors() -> None:
         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.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)
@@ -485,7 +487,7 @@ def test_series_dt_accessors() -> None:
         pd.Series([pd.Timedelta("1 day"), pd.Timedelta("2 days")]),
     )
 
-    check(assert_type(s5.dt.unit, str), str)
+    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)
@@ -557,7 +559,7 @@ 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, str), str)
+    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)
@@ -582,7 +584,7 @@ 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, str), str)
+    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)