Skip to content

Commit e29f129

Browse files
authored
Add accessors for TimedeltaIndex (#309)
* WIP: TimedeltaIndex accessors * remove dtl to make mypy happy
1 parent 3cb14af commit e29f129

File tree

7 files changed

+85
-17
lines changed

7 files changed

+85
-17
lines changed

pandas-stubs/core/arrays/datetimes.pyi

+6-2
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,17 @@
11
from datetime import tzinfo
22

33
import numpy as np
4-
from pandas.core.arrays import datetimelike as dtl
4+
from pandas.core.arrays.datetimelike import (
5+
DatelikeOps,
6+
DatetimeLikeArrayMixin,
7+
TimelikeOps,
8+
)
59

610
from pandas.core.dtypes.dtypes import DatetimeTZDtype as DatetimeTZDtype
711

812
def tz_to_dtype(tz): ...
913

10-
class DatetimeArray(dtl.DatetimeLikeArrayMixin, dtl.TimelikeOps, dtl.DatelikeOps):
14+
class DatetimeArray(DatetimeLikeArrayMixin, TimelikeOps, DatelikeOps):
1115
__array_priority__: int = ...
1216
def __init__(self, values, dtype=..., freq=..., copy: bool = ...) -> None: ...
1317
# ignore in dtype() is from the pandas source

pandas-stubs/core/arrays/period.pyi

+5-2
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,17 @@
11
from typing import Sequence
22

33
import numpy as np
4-
from pandas.core.arrays import datetimelike as dtl
4+
from pandas.core.arrays.datetimelike import (
5+
DatelikeOps,
6+
DatetimeLikeArrayMixin,
7+
)
58

69
from pandas._libs.tslibs import Timestamp
710
from pandas._libs.tslibs.period import Period as Period
811

912
from pandas.tseries.offsets import Tick as Tick
1013

11-
class PeriodArray(dtl.DatetimeLikeArrayMixin, dtl.DatelikeOps):
14+
class PeriodArray(DatetimeLikeArrayMixin, DatelikeOps):
1215
__array_priority__: int = ...
1316
def __init__(self, values, freq=..., dtype=..., copy: bool = ...) -> None: ...
1417
def dtype(self): ...

pandas-stubs/core/arrays/timedeltas.pyi

+5-2
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
from datetime import timedelta
22
from typing import Sequence
33

4-
from pandas.core.arrays import datetimelike as dtl
4+
from pandas.core.arrays.datetimelike import (
5+
DatetimeLikeArrayMixin,
6+
TimelikeOps,
7+
)
58

6-
class TimedeltaArray(dtl.DatetimeLikeArrayMixin, dtl.TimelikeOps):
9+
class TimedeltaArray(DatetimeLikeArrayMixin, TimelikeOps):
710
__array_priority__: int = ...
811
@property
912
def dtype(self): ...

pandas-stubs/core/indexes/accessors.pyi

+46-8
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ from pandas import (
1313
Index,
1414
PeriodIndex,
1515
Timedelta,
16+
TimedeltaIndex,
1617
)
1718
from pandas.core.accessor import PandasDelegate
1819
from pandas.core.arrays import (
@@ -156,6 +157,7 @@ _DTRoundingMethodReturnType = TypeVar(
156157
TimedeltaSeries,
157158
TimestampSeries,
158159
DatetimeIndex,
160+
# TimedeltaIndex
159161
)
160162

161163
class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]):
@@ -278,23 +280,30 @@ class DatetimeProperties(
278280
def to_pydatetime(self) -> np.ndarray: ...
279281
def isocalendar(self) -> DataFrame: ...
280282

281-
class _TimedeltaPropertiesNoRounding:
283+
_TDNoRoundingMethodReturnType = TypeVar(
284+
"_TDNoRoundingMethodReturnType", Series[int], Index
285+
)
286+
_TDTotalSecondsReturnType = TypeVar("_TDTotalSecondsReturnType", Series[float], Index)
287+
288+
class _TimedeltaPropertiesNoRounding(
289+
Generic[_TDNoRoundingMethodReturnType, _TDTotalSecondsReturnType]
290+
):
282291
def to_pytimedelta(self) -> np.ndarray: ...
283292
@property
284293
def components(self) -> DataFrame: ...
285294
@property
286-
def days(self) -> Series[int]: ...
295+
def days(self) -> _TDNoRoundingMethodReturnType: ...
287296
@property
288-
def seconds(self) -> Series[int]: ...
297+
def seconds(self) -> _TDNoRoundingMethodReturnType: ...
289298
@property
290-
def microseconds(self) -> Series[int]: ...
299+
def microseconds(self) -> _TDNoRoundingMethodReturnType: ...
291300
@property
292-
def nanoseconds(self) -> Series[int]: ...
293-
def total_seconds(self) -> Series[float]: ...
301+
def nanoseconds(self) -> _TDNoRoundingMethodReturnType: ...
302+
def total_seconds(self) -> _TDTotalSecondsReturnType: ...
294303

295304
class TimedeltaProperties(
296305
Properties,
297-
_TimedeltaPropertiesNoRounding,
306+
_TimedeltaPropertiesNoRounding[Series[int], Series[float]],
298307
_DatetimeRoundingMethods[TimedeltaSeries],
299308
): ...
300309

@@ -337,7 +346,7 @@ class CombinedDatetimelikeProperties(
337346
Series[str],
338347
PeriodSeries,
339348
],
340-
_TimedeltaPropertiesNoRounding,
349+
_TimedeltaPropertiesNoRounding[Series[int], Series[float]],
341350
_PeriodProperties,
342351
):
343352
def __new__(cls, data: Series): ...
@@ -379,3 +388,32 @@ class DatetimeIndexProperties(
379388
def std(
380389
self, axis: int | None = ..., ddof: int = ..., skipna: bool = ...
381390
) -> Timedelta: ...
391+
392+
# For some reason, using TimedeltaIndex as an argument to _DatetimeRoundingMethods
393+
# doesn't work for pyright. So we just make the rounding methods explicit here.
394+
class TimedeltaIndexProperties(
395+
Properties,
396+
_TimedeltaPropertiesNoRounding[Index, Index],
397+
# _DatetimeRoundingMethods[TimedeltaIndex],
398+
):
399+
def round(
400+
self,
401+
freq: str | BaseOffset | None,
402+
ambiguous: Literal["raise", "infer", "NaT"] | np_ndarray_bool = ...,
403+
nonexistent: Literal["shift_forward", "shift_backward", "NaT", "raise"]
404+
| Timedelta = ...,
405+
) -> TimedeltaIndex: ...
406+
def floor(
407+
self,
408+
freq: str | BaseOffset | None,
409+
ambiguous: Literal["raise", "infer", "NaT"] | np_ndarray_bool = ...,
410+
nonexistent: Literal["shift_forward", "shift_backward", "NaT", "raise"]
411+
| Timedelta = ...,
412+
) -> TimedeltaIndex: ...
413+
def ceil(
414+
self,
415+
freq: str | BaseOffset | None,
416+
ambiguous: Literal["raise", "infer", "NaT"] | np_ndarray_bool = ...,
417+
nonexistent: Literal["shift_forward", "shift_backward", "NaT", "raise"]
418+
| Timedelta = ...,
419+
) -> TimedeltaIndex: ...

pandas-stubs/core/indexes/datetimes.pyi

+1-1
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,12 @@ import numpy as np
55
from pandas import (
66
DataFrame,
77
Timedelta,
8+
TimedeltaIndex,
89
Timestamp,
910
)
1011
from pandas.core.indexes.accessors import DatetimeIndexProperties
1112
from pandas.core.indexes.api import Float64Index
1213
from pandas.core.indexes.datetimelike import DatetimeTimedeltaMixin
13-
from pandas.core.indexes.timedeltas import TimedeltaIndex
1414
from pandas.core.series import (
1515
TimedeltaSeries,
1616
TimestampSeries,

pandas-stubs/core/indexes/timedeltas.pyi

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from typing import overload
22

3-
from pandas.core.arrays.datetimelike import TimelikeOps
3+
from pandas.core.indexes.accessors import TimedeltaIndexProperties
44
from pandas.core.indexes.datetimelike import DatetimeTimedeltaMixin
55
from pandas.core.indexes.datetimes import DatetimeIndex
66
from pandas.core.series import TimedeltaSeries
@@ -11,7 +11,7 @@ from pandas._libs import (
1111
)
1212
from pandas._typing import num
1313

14-
class TimedeltaIndex(DatetimeTimedeltaMixin, TimelikeOps):
14+
class TimedeltaIndex(DatetimeTimedeltaMixin, TimedeltaIndexProperties):
1515
def __new__(
1616
cls,
1717
data=...,

tests/test_timefuncs.py

+20
Original file line numberDiff line numberDiff line change
@@ -436,6 +436,26 @@ def test_datetimeindex_accessors() -> None:
436436
check(assert_type(i0.is_normalized, bool), bool)
437437

438438

439+
def test_timedeltaindex_accessors() -> None:
440+
# GH 292
441+
i0 = pd.date_range("1/1/2021", "1/5/2021") - pd.Timestamp("1/3/2019")
442+
check(assert_type(i0, pd.TimedeltaIndex), pd.TimedeltaIndex)
443+
check(assert_type(i0.days, pd.Index), pd.Index, int)
444+
check(assert_type(i0.seconds, pd.Index), pd.Index, int)
445+
check(assert_type(i0.microseconds, pd.Index), pd.Index, int)
446+
check(assert_type(i0.nanoseconds, pd.Index), pd.Index, int)
447+
check(assert_type(i0.components, pd.DataFrame), pd.DataFrame)
448+
check(assert_type(i0.to_pytimedelta(), np.ndarray), np.ndarray)
449+
check(assert_type(i0.total_seconds(), pd.Index), pd.Index, float)
450+
check(
451+
assert_type(i0.round("D"), pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta
452+
)
453+
check(
454+
assert_type(i0.floor("D"), pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta
455+
)
456+
check(assert_type(i0.ceil("D"), pd.TimedeltaIndex), pd.TimedeltaIndex, pd.Timedelta)
457+
458+
439459
def test_some_offsets() -> None:
440460
# GH 222
441461

0 commit comments

Comments
 (0)