Add unit property and as_unit method to Timestamp, Timedelta …

…and `NaTType` (#864) * Add unit property and as_unit method to NaTType * Add unit property to Timestamp * Add as_unit method to Timestamp * Add unit property to Timedelta * Add as_unit method to Timedelta * Add TimeUnit type alias * Use TimeUnit type instead of str or Literal["s", "ms", "us", "ns"]
pandas-dev · Feb 13, 2024 · fdbe1df · fdbe1df
1 parent 7978238
commit fdbe1df
Show file tree

Hide file tree

Showing 9 changed files with 51 additions and 19 deletions.
diff --git a/pandas-stubs/_libs/tslibs/nattype.pyi b/pandas-stubs/_libs/tslibs/nattype.pyi
@@ -5,9 +5,13 @@ from datetime import (
 )
 
 import numpy as np
-from typing_extensions import TypeAlias
+from typing_extensions import (
+    Self,
+    TypeAlias,
+)
 
 from pandas._libs.tslibs.period import Period
+from pandas._typing import TimeUnit
 
 NaT: NaTType
 iNaT: int
@@ -121,3 +125,6 @@ class NaTType:
     __le__: _NatComparison
     __gt__: _NatComparison
     __ge__: _NatComparison
+    @property
+    def unit(self) -> TimeUnit: ...
+    def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...
diff --git a/pandas-stubs/_libs/tslibs/timedeltas.pyi b/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
@@ -390,3 +393,6 @@ class Timedelta(timedelta):
     @property
     def components(self) -> Components: ...
     def view(self, dtype: npt.DTypeLike = ...) -> object: ...
+    @property
+    def unit(self) -> TimeUnit: ...
+    def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...
diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/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,
 )
@@ -310,3 +311,6 @@ class Timestamp(datetime):
     def days_in_month(self) -> int: ...
     @property
     def daysinmonth(self) -> int: ...
+    @property
+    def unit(self) -> TimeUnit: ...
+    def as_unit(self, unit: TimeUnit, round_ok: bool = ...) -> Self: ...
diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi
@@ -775,6 +775,7 @@ RandomState: TypeAlias = (
     | np.random.RandomState
 )
 Frequency: TypeAlias = str | BaseOffset
+TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"]
 TimeGrouperOrigin: TypeAlias = (
     Timestamp | Literal["epoch", "start", "start_day", "end", "end_day"]
 )

diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi
@@ -116,6 +116,7 @@ from pandas._typing import (
     Suffixes,
     T as _T,
     TimestampConvention,
+    TimeUnit,
     ValidationOptions,
     WriteBuffer,
     XMLParsers,
@@ -2090,7 +2091,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
         ) = ...,
@@ -2109,7 +2110,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
         ) = ...,
@@ -2127,7 +2128,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
         ) = ...,
@@ -2145,7 +2146,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
         ) = ...,

diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/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(

diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi
@@ -145,6 +145,7 @@ from pandas._typing import (
     TimedeltaDtypeArg,
     TimestampConvention,
     TimestampDtypeArg,
+    TimeUnit,
     UIntDtypeArg,
     VoidDtypeArg,
     WriteBuffer,
@@ -493,7 +494,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
         ) = ...,
@@ -512,7 +513,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
         ) = ...,
@@ -530,7 +531,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
         ) = ...,
@@ -548,7 +549,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
         ) = ...,

diff --git a/pandas-stubs/io/json/_json.pyi b/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"]

diff --git a/tests/test_scalars.py b/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,6 +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, TimeUnit), str)
 
     check(assert_type(td.ceil("D"), pd.Timedelta), pd.Timedelta)
     check(assert_type(td.floor(Day()), pd.Timedelta), pd.Timedelta)
@@ -529,6 +531,11 @@ def test_timedelta_properties_methods() -> None:
     check(assert_type(td.view(np.int64), object), np.int64)
     check(assert_type(td.view("i8"), object), np.int64)
 
+    check(assert_type(td.as_unit("s"), pd.Timedelta), pd.Timedelta)
+    check(assert_type(td.as_unit("ms"), pd.Timedelta), pd.Timedelta)
+    check(assert_type(td.as_unit("us", round_ok=True), pd.Timedelta), pd.Timedelta)
+    check(assert_type(td.as_unit("ns", round_ok=False), pd.Timedelta), pd.Timedelta)
+
 
 def test_timedelta_add_sub() -> None:
     td = pd.Timedelta("1 day")
@@ -1189,6 +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, TimeUnit), str)
 
 
 def test_timestamp_add_sub() -> None:
@@ -1645,6 +1653,11 @@ def test_timestamp_misc_methods() -> None:
             pd.Timestamp,
         )
 
+    check(assert_type(ts2.as_unit("s"), pd.Timestamp), pd.Timestamp)
+    check(assert_type(ts2.as_unit("ms"), pd.Timestamp), pd.Timestamp)
+    check(assert_type(ts2.as_unit("us", round_ok=True), pd.Timestamp), pd.Timestamp)
+    check(assert_type(ts2.as_unit("ns", round_ok=False), pd.Timestamp), pd.Timestamp)
+
 
 def test_timestamp_types_arithmetic() -> None:
     ts: pd.Timestamp = pd.to_datetime("2021-03-01")