diff --git a/pandas-stubs/_libs/tslibs/nattype.pyi b/pandas-stubs/_libs/tslibs/nattype.pyi index 0fc2ba3d2..cc74db157 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 51dedb581..13fc53b8d 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 a8bf36b87..af856ae6f 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 89a7cbe72..9258054e3 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 faae35994..d64d1cee7 100644 --- a/pandas-stubs/core/frame.pyi +++ b/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 ) = ..., diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi index 50c2778ad..6d23fe56a 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 383c8df99..45fb2b20e 100644 --- a/pandas-stubs/core/series.pyi +++ b/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 ) = ..., diff --git a/pandas-stubs/io/json/_json.pyi b/pandas-stubs/io/json/_json.pyi index ea1fcb5bd..8151a6459 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 478aab2ee..dc6e0e5e4 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")