Skip to content

Commit 1e1b828

Browse files
Dr-Irvgramster
Dr-Irv
and
gramster
authored
pandas: get types of Interval correct (#167)
* get types of Interval correct * remove datetime * fix operators. use timestamp and timedelta from pandas * remove new use init * further refinement - now works with pandas mypy check * feedback on pandas PR * feedback from pandas core team review * remove generic from IntervalTree Co-authored-by: Graham Wheeler <[email protected]>
1 parent c2b5349 commit 1e1b828

File tree

4 files changed

+436
-257
lines changed

4 files changed

+436
-257
lines changed

‎partial/pandas/_libs/interval.pyi

Lines changed: 154 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,71 +1,175 @@
11
from __future__ import annotations
2-
import sys
3-
from typing import Generic, overload, Union, Protocol, TypeVar
4-
from _typing import Timedelta, Timestamp
5-
import datetime
62

7-
if sys.version_info >= (3, 8):
8-
from typing import Literal
9-
else:
10-
from typing_extensions import Literal
3+
from typing import (
4+
Any,
5+
Generic,
6+
Literal,
7+
Tuple,
8+
TypeVar,
9+
Union,
10+
overload,
11+
)
1112

12-
OrderableScalar = TypeVar("OrderableScalar", int, float)
13-
OrderableTimes = TypeVar("OrderableTimes", datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta)
14-
Orderable = TypeVar("Orderable", int, float, datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta)
13+
import numpy as np
14+
import numpy.typing as npt
1515

16-
class IntervalMixinProtocol(Protocol): ...
16+
from pandas._typing import (
17+
Timedelta,
18+
Timestamp,
19+
)
20+
21+
VALID_CLOSED: frozenset[str]
22+
23+
_OrderableScalarT = TypeVar("_OrderableScalarT", int, float)
24+
_OrderableTimesT = TypeVar("_OrderableTimesT", Timestamp, Timedelta)
25+
_OrderableT = TypeVar("_OrderableT", int, float, Timestamp, Timedelta)
26+
27+
class _LengthDescriptor:
28+
@overload
29+
def __get__(self, instance: Interval[float], owner: Any) -> float: ...
30+
@overload
31+
def __get__(self, instance: Interval[int], owner: Any) -> int: ...
32+
@overload
33+
def __get__(
34+
self, instance: Interval[_OrderableTimesT], owner: Any
35+
) -> Timedelta: ...
36+
@overload
37+
def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ...
38+
39+
class _MidDescriptor:
40+
@overload
41+
def __get__(self, instance: Interval[_OrderableScalarT], owner: Any) -> float: ...
42+
@overload
43+
def __get__(
44+
self, instance: Interval[_OrderableTimesT], owner: Any
45+
) -> _OrderableTimesT: ...
46+
@overload
47+
def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ...
1748

1849
class IntervalMixin:
1950
@property
20-
def closed_left(self: IntervalMixinProtocol) -> bool: ...
51+
def closed_left(self) -> bool: ...
2152
@property
22-
def closed_right(self: IntervalMixinProtocol) -> bool: ...
53+
def closed_right(self) -> bool: ...
2354
@property
24-
def open_left(self: IntervalMixinProtocol) -> bool: ...
55+
def open_left(self) -> bool: ...
2556
@property
26-
def open_right(self: IntervalMixinProtocol) -> bool: ...
57+
def open_right(self) -> bool: ...
58+
mid: _MidDescriptor
59+
length: _LengthDescriptor
2760
@property
28-
def mid(self: IntervalMixinProtocol) -> float: ...
29-
@property
30-
def is_empty(self: IntervalMixinProtocol) -> bool: ...
61+
def is_empty(self) -> bool: ...
62+
def _check_closed_matches(self, other: IntervalMixin, name: str = ...) -> None: ...
3163

32-
class Interval(IntervalMixin, Generic[Orderable]):
33-
@overload
34-
def left(self: Interval[OrderableScalar]) -> OrderableScalar: ...
35-
@overload
36-
def left(self: Interval[OrderableTimes]) -> OrderableTimes: ...
64+
class Interval(IntervalMixin, Generic[_OrderableT]):
3765
@property
38-
@overload
39-
def left(self: Interval[Orderable]) -> Orderable: ...
40-
@overload
41-
def right(self: Interval[OrderableScalar]) -> OrderableScalar: ...
42-
@overload
43-
def right(self: Interval[OrderableTimes]) -> OrderableTimes: ...
66+
def left(self: Interval[_OrderableT]) -> _OrderableT: ...
4467
@property
45-
@overload
46-
def right(self: Interval[Orderable]) -> Orderable: ...
68+
def right(self: Interval[_OrderableT]) -> _OrderableT: ...
4769
@property
48-
def closed(self) -> str: ...
70+
def closed(self) -> Literal["left", "right", "both", "neither"]: ...
4971
def __init__(
5072
self,
51-
left: Orderable,
52-
right: Orderable,
53-
closed: Union[str, Literal["left", "right", "both", "neither"]] = ...,
54-
) -> None: ...
55-
@overload
56-
def length(self: Interval[OrderableScalar]) -> float: ...
73+
left: _OrderableT,
74+
right: _OrderableT,
75+
closed: Literal["left", "right", "both", "neither"] = ...,
76+
): ...
77+
def __hash__(self) -> int: ...
5778
@overload
58-
def length(self: Interval[OrderableTimes]) -> Timedelta: ...
59-
@property
79+
def __contains__(self: Interval[_OrderableTimesT], _OrderableTimesT) -> bool: ...
6080
@overload
61-
def length(self: Interval[Orderable]) -> Orderable: ...
62-
def __hash__(self) -> int: ...
63-
def __contains__(self, key: Orderable) -> bool: ...
81+
def __contains__(
82+
self: Interval[_OrderableScalarT], key: Union[int, float]
83+
) -> bool: ...
6484
def __repr__(self) -> str: ...
6585
def __str__(self) -> str: ...
66-
def __add__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ...
67-
def __sub__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ...
68-
def __mul__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ...
69-
def __truediv__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ...
70-
def __floordiv__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ...
71-
def overlaps(self: Interval[Orderable], other: Interval[Orderable]) -> bool: ...
86+
@overload
87+
def __add__(
88+
self: Interval[_OrderableTimesT], y: Timedelta
89+
) -> Interval[_OrderableTimesT]: ...
90+
@overload
91+
def __add__(self: Interval[int], y: int) -> Interval[int]: ...
92+
@overload
93+
def __add__(self: Interval[int], y: float) -> Interval[float]: ...
94+
@overload
95+
def __add__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ...
96+
@overload
97+
def __radd__(
98+
self: Interval[_OrderableTimesT], y: Timedelta
99+
) -> Interval[_OrderableTimesT]: ...
100+
@overload
101+
def __radd__(self: Interval[int], y: int) -> Interval[int]: ...
102+
@overload
103+
def __radd__(self: Interval[int], y: float) -> Interval[float]: ...
104+
@overload
105+
def __radd__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ...
106+
@overload
107+
def __sub__(
108+
self: Interval[_OrderableTimesT], y: Timedelta
109+
) -> Interval[_OrderableTimesT]: ...
110+
@overload
111+
def __sub__(self: Interval[int], y: int) -> Interval[int]: ...
112+
@overload
113+
def __sub__(self: Interval[int], y: float) -> Interval[float]: ...
114+
@overload
115+
def __sub__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ...
116+
@overload
117+
def __rsub__(
118+
self: Interval[_OrderableTimesT], y: Timedelta
119+
) -> Interval[_OrderableTimesT]: ...
120+
@overload
121+
def __rsub__(self: Interval[int], y: int) -> Interval[int]: ...
122+
@overload
123+
def __rsub__(self: Interval[int], y: float) -> Interval[float]: ...
124+
@overload
125+
def __rsub__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ...
126+
@overload
127+
def __mul__(self: Interval[int], y: int) -> Interval[int]: ...
128+
@overload
129+
def __mul__(self: Interval[int], y: float) -> Interval[float]: ...
130+
@overload
131+
def __mul__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ...
132+
@overload
133+
def __rmul__(self: Interval[int], y: int) -> Interval[int]: ...
134+
@overload
135+
def __rmul__(self: Interval[int], y: float) -> Interval[float]: ...
136+
@overload
137+
def __rmul__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ...
138+
@overload
139+
def __truediv__(self: Interval[int], y: int) -> Interval[int]: ...
140+
@overload
141+
def __truediv__(self: Interval[int], y: float) -> Interval[float]: ...
142+
@overload
143+
def __truediv__(self: Interval[float], y: Union[int, float]) -> Interval[float]: ...
144+
@overload
145+
def __floordiv__(self: Interval[int], y: int) -> Interval[int]: ...
146+
@overload
147+
def __floordiv__(self: Interval[int], y: float) -> Interval[float]: ...
148+
@overload
149+
def __floordiv__(
150+
self: Interval[float], y: Union[int, float]
151+
) -> Interval[float]: ...
152+
def overlaps(self: Interval[_OrderableT], other: Interval[_OrderableT]) -> bool: ...
153+
154+
def intervals_to_interval_bounds(
155+
intervals: np.ndarray, validate_closed: bool = ...
156+
) -> Tuple[np.ndarray, np.ndarray, str]: ...
157+
158+
class IntervalTree(IntervalMixin):
159+
def __init__(
160+
self,
161+
left: np.ndarray,
162+
right: np.ndarray,
163+
closed: Literal["left", "right", "both", "neither"] = ...,
164+
leaf_size: int = ...,
165+
): ...
166+
def get_indexer(self, target) -> npt.NDArray[np.intp]: ...
167+
def get_indexer_non_unique(
168+
self, target
169+
) -> Tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ...
170+
_na_count: int
171+
@property
172+
def is_overlapping(self) -> bool: ...
173+
@property
174+
def is_monotonic_increasing(self) -> bool: ...
175+
def clear_mapping(self) -> None: ...

‎partial/pandas/_libs/tslibs/timedeltas.pyi

Lines changed: 67 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -1,84 +1,83 @@
1-
from __future__ import annotations
21
from datetime import timedelta
3-
from typing import Any, Union, Tuple, Type, Optional, Sequence
2+
from typing import (
3+
ClassVar,
4+
Type,
5+
TypeVar,
6+
overload,
7+
)
48

59
import numpy as np
6-
from pandas._typing import Dtype
710

8-
class _Timedelta(timedelta):
9-
def __hash__(self) -> int: ...
10-
def __richcmp__(self, other, op: int) -> Any: ...
11-
def to_timedelta64(self) -> np.timedelta64: ...
12-
def view(self, dtype: Dtype): ...
13-
@property
14-
def components(self) -> Tuple: ... # Really a namedtuple
15-
@property
16-
def delta(self) -> int: ...
17-
@property
18-
def asm8(self) -> np.timedelta64: ...
19-
@property
20-
def resolution_string(self) -> str: ...
21-
@property
22-
def nanoseconds(self) -> int: ...
23-
def __repr__(self) -> str: ...
24-
def __str__(self) -> str: ...
25-
def __bool__(self) -> bool: ...
26-
def isoformat(self) -> str: ...
11+
from pandas._libs.tslibs import (
12+
NaTType,
13+
Tick,
14+
)
15+
from pandas._typing import npt
2716

28-
class Timedelta(_Timedelta):
29-
def __new__(cls, value: object = ..., unit: Optional[str] = ..., **kwargs) -> Timedelta: ...
30-
def __setstate__(self, state) -> None: ...
31-
def __reduce__(self) -> Tuple[Type[Timedelta], int]: ...
32-
def round(self, freq: str) -> Timedelta: ...
33-
def floor(self, freq: str) -> Timedelta: ...
34-
def ceil(self, freq: str) -> Timedelta: ...
35-
def __inv__(self) -> Timedelta: ...
36-
def __neg__(self) -> Timedelta: ...
37-
def __pos__(self) -> Timedelta: ...
38-
def __abs__(self) -> Timedelta: ...
39-
def __add__(self, other) -> Timedelta: ...
40-
def __radd__(self, other) -> Timedelta: ...
41-
def __sub__(self, other) -> Timedelta: ...
42-
def __rsub_(self, other) -> Timedelta: ...
43-
def __mul__(self, other) -> Timedelta: ...
44-
__rmul__ = __mul__
45-
def __truediv__(self, other) -> Union[float, Timedelta]: ...
46-
def __rtruediv__(self, other) -> float: ...
47-
def __floordiv__(self, other) -> Union[int, Timedelta]: ...
48-
def __rfloordiv__(self, other) -> int: ...
49-
def __mod__(self, other) -> int: ...
50-
def __rmod__(self, other) -> int: ...
51-
def __divmod__(self, other) -> Tuple[int, int]: ...
52-
def __rdivmod__(self, other) -> Tuple[int, int]: ...
17+
_S = TypeVar("_S", bound=timedelta)
5318

54-
@property
55-
def asm8(self) -> int: ...
56-
@property
57-
def components(self) -> int: ...
19+
def ints_to_pytimedelta(arr: Sequence[int], box: bool = ...) -> np.ndarray: ...
20+
def array_to_timedelta64(
21+
values: npt.NDArray[np.object_],
22+
unit: str | None = ...,
23+
errors: str = ...,
24+
) -> np.ndarray: ... # np.ndarray[m8ns]
25+
def parse_timedelta_unit(unit: str | None) -> str: ...
26+
def delta_to_nanoseconds(delta: Tick | np.timedelta64 | timedelta | int) -> int: ...
27+
28+
class Timedelta(timedelta):
29+
min: ClassVar[Timedelta]
30+
max: ClassVar[Timedelta]
31+
resolution: ClassVar[Timedelta]
32+
value: int # np.int64
33+
34+
# error: "__new__" must return a class instance (got "Union[Timedelta, NaTType]")
35+
def __new__( # type: ignore[misc]
36+
cls: Type[_S],
37+
value=...,
38+
unit: str = ...,
39+
**kwargs: int | float | np.integer | np.floating,
40+
) -> _S | NaTType: ...
5841
@property
5942
def days(self) -> int: ...
6043
@property
61-
def delta(self) -> int: ...
44+
def seconds(self) -> int: ...
6245
@property
6346
def microseconds(self) -> int: ...
47+
def total_seconds(self) -> float: ...
48+
def to_pytimedelta(self) -> timedelta: ...
49+
def to_timedelta64(self) -> np.timedelta64: ...
6450
@property
65-
def nanoseconds(self) -> int: ...
66-
@property
67-
def resolution_string(self) -> str: ...
68-
@property
69-
def seconds(self) -> int: ...
70-
max: Timedelta = ...
71-
min: Timedelta = ...
72-
resolution: 'Timedelta' = ...
51+
def asm8(self) -> np.timedelta64: ...
52+
# TODO: round/floor/ceil could return NaT?
53+
def round(self, freq) -> Timedelta: ...
7354
def ceil(self, freq, **kwargs) -> Timedelta: ...
7455
def floor(self, freq, **kwargs) -> Timedelta: ...
56+
@property
57+
def resolution_string(self) -> str: ...
58+
def __add__(self, other: timedelta) -> timedelta: ...
59+
def __radd__(self, other: timedelta) -> timedelta: ...
60+
def __sub__(self, other: timedelta) -> timedelta: ...
61+
def __rsub__(self, other: timedelta) -> timedelta: ...
62+
def __neg__(self) -> timedelta: ...
63+
def __pos__(self) -> timedelta: ...
64+
def __abs__(self) -> timedelta: ...
65+
def __mul__(self, other: float) -> timedelta: ...
66+
def __rmul__(self, other: float) -> timedelta: ...
67+
@overload
68+
def __floordiv__(self, other: timedelta) -> int: ...
69+
@overload
70+
def __floordiv__(self, other: int) -> timedelta: ...
71+
@overload
72+
def __truediv__(self, other: timedelta) -> float: ...
73+
@overload
74+
def __truediv__(self, other: float) -> timedelta: ...
75+
def __mod__(self, other: timedelta) -> timedelta: ...
76+
def __divmod__(self, other: timedelta) -> tuple[int, timedelta]: ...
77+
def __le__(self, other: timedelta) -> bool: ...
78+
def __lt__(self, other: timedelta) -> bool: ...
79+
def __ge__(self, other: timedelta) -> bool: ...
80+
def __gt__(self, other: timedelta) -> bool: ...
81+
def __hash__(self) -> int: ...
7582
def isoformat(self) -> str: ...
76-
def round(self, freq) -> Timedelta: ...
7783
def to_numpy(self) -> _np.timedelta64: ...
78-
def to_pytimedelta(self) -> datetime.timedelta: ...
79-
def to_timedelta64(self) -> _np.timedelta64: ...
80-
def total_seconds(self) -> float: ...
81-
82-
83-
def delta_to_nanoseconds(delta: _Timedelta) -> int: ...
84-
def ints_to_pytimedelta(arr: Sequence[int], box: bool = ...) -> np.ndarray: ...

0 commit comments

Comments
 (0)