From fdbe1df6128fe7042398a722062588c51e331932 Mon Sep 17 00:00:00 2001 From: Soshi Katsuta Date: Tue, 13 Feb 2024 23:40:06 +0900 Subject: [PATCH] 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-stubs/_libs/tslibs/nattype.pyi | 9 ++++++++- pandas-stubs/_libs/tslibs/timedeltas.pyi | 8 +++++++- pandas-stubs/_libs/tslibs/timestamps.pyi | 4 ++++ pandas-stubs/_typing.pyi | 1 + pandas-stubs/core/frame.pyi | 9 +++++---- pandas-stubs/core/indexes/datetimes.pyi | 8 +++----- pandas-stubs/core/series.pyi | 9 +++++---- pandas-stubs/io/json/_json.pyi | 9 +++++---- tests/test_scalars.py | 13 +++++++++++++ 9 files changed, 51 insertions(+), 19 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/nattype.pyi b/pandas-stubs/_libs/tslibs/nattype.pyi index 0fc2ba3d..cc74db15 100644 --- 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 index 51dedb58..13fc53b8 100644 --- 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 index a8bf36b8..af856ae6 100644 --- 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 index 89a7cbe7..9258054e 100644 --- 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 index 8656099c..7da14910 100644 --- 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 index 50c2778a..6d23fe56 100644 --- 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 index d8c691a3..46a8ce14 100644 --- 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 index ea1fcb5b..8151a645 100644 --- 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 index 478aab2e..dc6e0e5e 100644 --- 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")