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

Author: skatsuta

pandas-stubs/_libs/tslibs/nattype.pyi

@@ -3,7 +3,6 @@ from datetime import (
     timedelta,
     tzinfo as _tzinfo,
 )
-from typing import Literal
 
 import numpy as np
 from typing_extensions import (
@@ -12,6 +11,7 @@ from typing_extensions import (
 )
 
 from pandas._libs.tslibs.period import Period
+from pandas._typing import TimeUnit
 
 NaT: NaTType
 iNaT: int
@@ -126,7 +126,5 @@ class NaTType:
     __gt__: _NatComparison
     __ge__: _NatComparison
     @property
-    def unit(self) -> str: ...
-    def as_unit(
-        self, unit: Literal["s", "ms", "us", "ns"], round_ok: bool = ...
-    ) -> Self: ...
+    def unit(self) -> TimeUnit: ...
+    def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...

pandas-stubs/_libs/tslibs/timedeltas.pyi

@@ -31,7 +31,10 @@ from pandas._libs.tslibs import (
 )
 from pandas._libs.tslibs.period import Period
 from pandas._libs.tslibs.timestamps import Timestamp
-from pandas._typing import npt
+from pandas._typing import (
+    TimeUnit,
+    npt,
+)
 
 class Components(NamedTuple):
     days: int
@@ -391,7 +394,5 @@ class Timedelta(timedelta):
     def components(self) -> Components: ...
     def view(self, dtype: npt.DTypeLike = ...) -> object: ...
     @property
-    def unit(self) -> str: ...
-    def as_unit(
-        self, unit: Literal["s", "ms", "us", "ns"], round_ok: bool = ...
-    ) -> Self: ...
+    def unit(self) -> TimeUnit: ...
+    def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...

pandas-stubs/_libs/tslibs/timestamps.pyi

@@ -38,6 +38,7 @@ from pandas._libs.tslibs import (
     Timedelta,
 )
 from pandas._typing import (
+    TimeUnit,
     np_ndarray_bool,
     npt,
 )
@@ -311,7 +312,5 @@ class Timestamp(datetime):
     @property
     def daysinmonth(self) -> int: ...
     @property
-    def unit(self) -> str: ...
-    def as_unit(
-        self, unit: Literal["s", "ms", "us", "ns"], round_ok: bool = ...
-    ) -> Self: ...
+    def unit(self) -> TimeUnit: ...
+    def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...

pandas-stubs/core/frame.pyi

@@ -115,6 +115,7 @@ from pandas._typing import (
     StrLike,
     Suffixes,
     TimestampConvention,
+    TimeUnit,
     ValidationOptions,
     WriteBuffer,
     XMLParsers,
@@ -2082,7 +2083,7 @@ class DataFrame(NDFrame, OpsMixin):
         date_format: Literal["epoch", "iso"] | None = ...,
         double_precision: int = ...,
         force_ascii: _bool = ...,
-        date_unit: Literal["s", "ms", "us", "ns"] = ...,
+        date_unit: TimeUnit = ...,
         default_handler: (
             Callable[[Any], _str | float | _bool | list | dict] | None
         ) = ...,
@@ -2101,7 +2102,7 @@ class DataFrame(NDFrame, OpsMixin):
         date_format: Literal["epoch", "iso"] | None = ...,
         double_precision: int = ...,
         force_ascii: _bool = ...,
-        date_unit: Literal["s", "ms", "us", "ns"] = ...,
+        date_unit: TimeUnit = ...,
         default_handler: (
             Callable[[Any], _str | float | _bool | list | dict] | None
         ) = ...,
@@ -2119,7 +2120,7 @@ class DataFrame(NDFrame, OpsMixin):
         date_format: Literal["epoch", "iso"] | None = ...,
         double_precision: int = ...,
         force_ascii: _bool = ...,
-        date_unit: Literal["s", "ms", "us", "ns"] = ...,
+        date_unit: TimeUnit = ...,
         default_handler: (
             Callable[[Any], _str | float | _bool | list | dict] | None
         ) = ...,
@@ -2137,7 +2138,7 @@ class DataFrame(NDFrame, OpsMixin):
         date_format: Literal["epoch", "iso"] | None = ...,
         double_precision: int = ...,
         force_ascii: _bool = ...,
-        date_unit: Literal["s", "ms", "us", "ns"] = ...,
+        date_unit: TimeUnit = ...,
         default_handler: (
             Callable[[Any], _str | float | _bool | list | dict] | None
         ) = ...,

pandas-stubs/core/indexes/datetimes.pyi

@@ -7,10 +7,7 @@ from datetime import (
     timedelta,
     tzinfo,
 )
-from typing import (
-    Literal,
-    overload,
-)
+from typing import overload
 
 import numpy as np
 from pandas import (
@@ -32,6 +29,7 @@ from pandas._typing import (
     ArrayLike,
     DateAndDatetimeLike,
     IntervalClosedType,
+    TimeUnit,
 )
 
 from pandas.core.dtypes.dtypes import DatetimeTZDtype
@@ -100,7 +98,7 @@ def date_range(
     normalize: bool = ...,
     name: Hashable | None = ...,
     inclusive: IntervalClosedType = ...,
-    unit: Literal["s", "ms", "us", "ns"] | None = ...,
+    unit: TimeUnit | None = ...,
 ) -> DatetimeIndex: ...
 @overload
 def bdate_range(

pandas-stubs/core/series.pyi

@@ -144,6 +144,7 @@ from pandas._typing import (
     TimedeltaDtypeArg,
     TimestampConvention,
     TimestampDtypeArg,
+    TimeUnit,
     UIntDtypeArg,
     VoidDtypeArg,
     WriteBuffer,
@@ -486,7 +487,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
         date_format: Literal["epoch", "iso"] | None = ...,
         double_precision: int = ...,
         force_ascii: _bool = ...,
-        date_unit: Literal["s", "ms", "us", "ns"] = ...,
+        date_unit: TimeUnit = ...,
         default_handler: (
             Callable[[Any], _str | float | _bool | list | dict] | None
         ) = ...,
@@ -505,7 +506,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
         date_format: Literal["epoch", "iso"] | None = ...,
         double_precision: int = ...,
         force_ascii: _bool = ...,
-        date_unit: Literal["s", "ms", "us", "ns"] = ...,
+        date_unit: TimeUnit = ...,
         default_handler: (
             Callable[[Any], _str | float | _bool | list | dict] | None
         ) = ...,
@@ -523,7 +524,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
         date_format: Literal["epoch", "iso"] | None = ...,
         double_precision: int = ...,
         force_ascii: _bool = ...,
-        date_unit: Literal["s", "ms", "us", "ns"] = ...,
+        date_unit: TimeUnit = ...,
         default_handler: (
             Callable[[Any], _str | float | _bool | list | dict] | None
         ) = ...,
@@ -541,7 +542,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
         date_format: Literal["epoch", "iso"] | None = ...,
         double_precision: int = ...,
         force_ascii: _bool = ...,
-        date_unit: Literal["s", "ms", "us", "ns"] = ...,
+        date_unit: TimeUnit = ...,
         default_handler: (
             Callable[[Any], _str | float | _bool | list | dict] | None
         ) = ...,

pandas-stubs/io/json/_json.pyi

@@ -22,6 +22,7 @@ from pandas._typing import (
     NDFrameT,
     ReadBuffer,
     StorageOptions,
+    TimeUnit,
 )
 
 @overload
@@ -35,7 +36,7 @@ def read_json(
     convert_dates: bool | list[str] = ...,
     keep_default_dates: bool = ...,
     precise_float: bool = ...,
-    date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
+    date_unit: TimeUnit | None = ...,
     encoding: str | None = ...,
     encoding_errors: (
         Literal["strict", "ignore", "replace", "backslashreplace", "surrogateescape"]
@@ -59,7 +60,7 @@ def read_json(
     convert_dates: bool | list[str] = ...,
     keep_default_dates: bool = ...,
     precise_float: bool = ...,
-    date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
+    date_unit: TimeUnit | None = ...,
     encoding: str | None = ...,
     encoding_errors: (
         Literal["strict", "ignore", "replace", "backslashreplace", "surrogateescape"]
@@ -83,7 +84,7 @@ def read_json(
     convert_dates: bool | list[str] = ...,
     keep_default_dates: bool = ...,
     precise_float: bool = ...,
-    date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
+    date_unit: TimeUnit | None = ...,
     encoding: str | None = ...,
     encoding_errors: (
         Literal["strict", "ignore", "replace", "backslashreplace", "surrogateescape"]
@@ -107,7 +108,7 @@ def read_json(
     convert_dates: bool | list[str] = ...,
     keep_default_dates: bool = ...,
     precise_float: bool = ...,
-    date_unit: Literal["s", "ms", "us", "ns"] | None = ...,
+    date_unit: TimeUnit | None = ...,
     encoding: str | None = ...,
     encoding_errors: (
         Literal["strict", "ignore", "replace", "backslashreplace", "surrogateescape"]

tests/test_scalars.py

@@ -24,6 +24,7 @@
     NaTType,
 )
 from pandas._libs.tslibs.timedeltas import Components
+from pandas._typing import TimeUnit
 
 from tests import (
     TYPE_CHECKING_INVALID_USAGE,
@@ -517,7 +518,7 @@ def test_timedelta_properties_methods() -> None:
     check(assert_type(td.value, int), int)
     check(assert_type(td.resolution_string, str), str)
     check(assert_type(td.components, Components), Components)
-    check(assert_type(td.unit, str), str)
+    check(assert_type(td.unit, TimeUnit), str)
 
     check(assert_type(td.ceil("D"), pd.Timedelta), pd.Timedelta)
     check(assert_type(td.floor(Day()), pd.Timedelta), pd.Timedelta)
@@ -1195,7 +1196,7 @@ def test_timestamp_properties() -> None:
     check(assert_type(ts.tzinfo, Optional[dt.tzinfo]), type(None))
     check(assert_type(ts.value, int), int)
     check(assert_type(ts.year, int), int)
-    check(assert_type(ts.unit, str), str)
+    check(assert_type(ts.unit, TimeUnit), str)
 
 
 def test_timestamp_add_sub() -> None: