Skip to content

Add unit property and as_unit method to DatetimeIndex, TimedeltaIndex and Series.dt #865

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Feb 13, 2024
Merged

Add unit property and as_unit method to DatetimeIndex, TimedeltaIndex and Series.dt #865

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5a37ab9
Add unit property and as_unit method to DatetimeIndex and TimedeltaIndex
skatsuta Feb 11, 2024
f571b68
Add unit property and as_unit method to TimelikeOps
skatsuta Feb 12, 2024
2ca04fa
Add unit property and as_unit method to TimestampSeries and Timedelta…
skatsuta Feb 13, 2024
9120dce
Use TimeUnit type alias instead of str or Literal["s", "ms", "us", "ns"]
skatsuta Feb 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions pandas-stubs/core/arrays/datetimelike.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,16 +5,21 @@ from pandas.core.arrays.base import (
ExtensionArray,
ExtensionOpsMixin,
)
from typing_extensions import Self

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) -> 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 = ...): ...
Expand Down
35 changes: 21 additions & 14 deletions pandas-stubs/core/indexes/accessors.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down Expand Up @@ -154,16 +155,16 @@ class _DatetimeLikeOps(
# type of the series, we don't know which kind of series was ...ed
# in to the dt accessor

_DTRoundingMethodReturnType = TypeVar(
"_DTRoundingMethodReturnType",
_DTTimestampTimedeltaReturnType = TypeVar(
"_DTTimestampTimedeltaReturnType",
Comment on lines -157 to +159
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rename to a general name to reuse this type var in the return type annotation for as_unit method

Series,
TimedeltaSeries,
TimestampSeries,
TimedeltaSeries,
DatetimeIndex,
TimedeltaIndex,
)

class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]):
class _DatetimeRoundingMethods(Generic[_DTTimestampTimedeltaReturnType]):
def round(
self,
freq: str | BaseOffset | None,
Expand All @@ -173,7 +174,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]):
| timedelta
| Timedelta
) = ...,
) -> _DTRoundingMethodReturnType: ...
) -> _DTTimestampTimedeltaReturnType: ...
def floor(
self,
freq: str | BaseOffset | None,
Expand All @@ -183,7 +184,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]):
| timedelta
| Timedelta
) = ...,
) -> _DTRoundingMethodReturnType: ...
) -> _DTTimestampTimedeltaReturnType: ...
def ceil(
self,
freq: str | BaseOffset | None,
Expand All @@ -193,7 +194,7 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]):
| timedelta
| Timedelta
) = ...,
) -> _DTRoundingMethodReturnType: ...
) -> _DTTimestampTimedeltaReturnType: ...

_DTNormalizeReturnType = TypeVar(
"_DTNormalizeReturnType", TimestampSeries, DatetimeIndex
Expand All @@ -202,9 +203,9 @@ _DTStrKindReturnType = TypeVar("_DTStrKindReturnType", Series[str], Index)
_DTToPeriodReturnType = TypeVar("_DTToPeriodReturnType", PeriodSeries, PeriodIndex)

class _DatetimeLikeNoTZMethods(
_DatetimeRoundingMethods[_DTRoundingMethodReturnType],
_DatetimeRoundingMethods[_DTTimestampTimedeltaReturnType],
Generic[
_DTRoundingMethodReturnType,
_DTTimestampTimedeltaReturnType,
_DTNormalizeReturnType,
_DTStrKindReturnType,
_DTToPeriodReturnType,
Expand Down Expand Up @@ -238,15 +239,15 @@ class _DatetimeNoTZProperties(
_DTFreqReturnType,
],
_DatetimeLikeNoTZMethods[
_DTRoundingMethodReturnType,
_DTTimestampTimedeltaReturnType,
_DTNormalizeReturnType,
_DTStrKindReturnType,
_DTToPeriodReturnType,
],
Generic[
_DTFieldOpsReturnType,
_DTBoolOpsReturnType,
_DTRoundingMethodReturnType,
_DTTimestampTimedeltaReturnType,
_DTOtherOpsDateReturnType,
_DTOtherOpsTimeReturnType,
_DTFreqReturnType,
Expand All @@ -261,7 +262,7 @@ class DatetimeProperties(
_DatetimeNoTZProperties[
_DTFieldOpsReturnType,
_DTBoolOpsReturnType,
_DTRoundingMethodReturnType,
_DTTimestampTimedeltaReturnType,
_DTOtherOpsDateReturnType,
_DTOtherOpsTimeReturnType,
_DTFreqReturnType,
Expand All @@ -272,7 +273,7 @@ class DatetimeProperties(
Generic[
_DTFieldOpsReturnType,
_DTBoolOpsReturnType,
_DTRoundingMethodReturnType,
_DTTimestampTimedeltaReturnType,
_DTOtherOpsDateReturnType,
_DTOtherOpsTimeReturnType,
_DTFreqReturnType,
Expand All @@ -283,6 +284,9 @@ class DatetimeProperties(
):
def to_pydatetime(self) -> np.ndarray: ...
def isocalendar(self) -> DataFrame: ...
@property
def unit(self) -> TimeUnit: ...
def as_unit(self, unit: TimeUnit) -> _DTTimestampTimedeltaReturnType: ...

_TDNoRoundingMethodReturnType = TypeVar(
"_TDNoRoundingMethodReturnType", Series[int], Index
Expand All @@ -309,7 +313,10 @@ class TimedeltaProperties(
Properties,
_TimedeltaPropertiesNoRounding[Series[int], Series[float]],
_DatetimeRoundingMethods[TimedeltaSeries],
): ...
):
@property
def unit(self) -> TimeUnit: ...
def as_unit(self, unit: TimeUnit) -> TimedeltaSeries: ...

_PeriodDTReturnTypes = TypeVar("_PeriodDTReturnTypes", TimestampSeries, DatetimeIndex)
_PeriodIntReturnTypes = TypeVar("_PeriodIntReturnTypes", Series[int], Index[int])
Expand Down
11 changes: 9 additions & 2 deletions pandas-stubs/core/indexes/datetimelike.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
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
Expand All @@ -19,4 +23,7 @@ class DatetimeIndexOpsMixin(ExtensionIndex[S1]):
self, other: DatetimeIndexOpsMixin
) -> TimedeltaIndex: ...

class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1]): ...
class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1]):
@property
def unit(self) -> TimeUnit: ...
def as_unit(self, unit: TimeUnit) -> Self: ...
47 changes: 47 additions & 0 deletions tests/test_timefuncs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
TYPE_CHECKING,
Any,
Optional,
cast,
)

import numpy as np
Expand All @@ -24,6 +25,8 @@
from pandas._typing import FulldatetimeDict
else:
FulldatetimeDict = Any
from pandas._typing import TimeUnit

from tests import (
TYPE_CHECKING_INVALID_USAGE,
check,
Expand Down Expand Up @@ -428,6 +431,11 @@ 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, 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)
check(assert_type(s0.dt.as_unit("ns"), "TimestampSeries"), pd.Series, pd.Timestamp)

i1 = pd.period_range(start="2022-06-01", periods=10)

Expand Down Expand Up @@ -455,6 +463,35 @@ 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, 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)
check(assert_type(s2.dt.as_unit("ns"), "TimedeltaSeries"), pd.Series, pd.Timedelta)

# Checks for general Series other than TimestampSeries and TimedeltaSeries

s4 = cast(
"pd.Series[pd.Timestamp]",
pd.Series([pd.Timestamp("2024-01-01"), pd.Timestamp("2024-01-02")]),
)

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)
check(assert_type(s4.dt.as_unit("ns"), pd.Series), pd.Series, pd.Timestamp)

s5 = cast(
"pd.Series[pd.Timedelta]",
pd.Series([pd.Timedelta("1 day"), pd.Timedelta("2 days")]),
)

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)
check(assert_type(s5.dt.as_unit("ns"), pd.Series), pd.Series, pd.Timedelta)


def test_datetimeindex_accessors() -> None:
Expand Down Expand Up @@ -522,6 +559,11 @@ 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, 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)
check(assert_type(i0.as_unit("ns"), pd.DatetimeIndex), pd.DatetimeIndex)


def test_timedeltaindex_accessors() -> None:
Expand All @@ -542,6 +584,11 @@ 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, 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)
check(assert_type(i0.as_unit("ns"), pd.TimedeltaIndex), pd.TimedeltaIndex)


def test_periodindex_accessors() -> None:
Expand Down