Skip to content

Commit 977f98f

Browse files
bashtagebashtage
committed
ENH: Improve typing for Interval
1 parent 5f86351 commit 977f98f

File tree

2 files changed

+169
-4
lines changed

2 files changed

+169
-4
lines changed

‎pandas-stubs/_libs/interval.pyi

+13-4
Original file line numberDiff line numberDiff line change
@@ -123,28 +123,37 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
123123
@overload
124124
def __mul__(self: Interval[float], y: float) -> Interval[float]: ...
125125
@overload
126+
def __mul__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
127+
@overload
126128
def __rmul__(
127129
self: Interval[int], y: _OrderableScalarT
128130
) -> Interval[_OrderableScalarT]: ...
129131
@overload
130132
def __rmul__(self: Interval[float], y: float) -> Interval[float]: ...
131133
@overload
134+
def __rmul__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
135+
@overload
132136
def __truediv__(
133137
self: Interval[int], y: _OrderableScalarT
134138
) -> Interval[_OrderableScalarT]: ...
135139
@overload
136140
def __truediv__(self: Interval[float], y: float) -> Interval[float]: ...
137141
@overload
142+
def __truediv__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
143+
@overload
138144
def __floordiv__(
139145
self: Interval[int], y: _OrderableScalarT
140146
) -> Interval[_OrderableScalarT]: ...
141147
@overload
142148
def __floordiv__(self: Interval[float], y: float) -> Interval[float]: ...
149+
@overload
150+
def __floordiv__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
151+
@overload
143152
def overlaps(self: Interval[_OrderableT], other: Interval[_OrderableT]) -> bool: ...
144-
145-
def intervals_to_interval_bounds(
146-
intervals: np.ndarray, validate_closed: bool = ...
147-
) -> tuple[np.ndarray, np.ndarray, str]: ...
153+
@overload
154+
def overlaps(self: Interval[int], other: Interval[float]) -> bool: ...
155+
@overload
156+
def overlaps(self: Interval[float], other: Interval[int]) -> bool: ...
148157

149158
class IntervalTree(IntervalMixin):
150159
def __init__(

‎tests/test_scalars.py

+156
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
from __future__ import annotations
2+
3+
from typing import Literal
4+
5+
import pandas as pd
6+
from typing_extensions import assert_type
7+
8+
from tests import check
9+
10+
11+
def test_interval() -> None:
12+
i0 = pd.Interval(0, 1, closed="left")
13+
i1 = pd.Interval(0.0, 1.0, closed="right")
14+
i2 = pd.Interval(
15+
pd.Timestamp("2017-01-01"), pd.Timestamp("2017-01-02"), closed="both"
16+
)
17+
i3 = pd.Interval(pd.Timedelta("1 days"), pd.Timedelta("2 days"), closed="neither")
18+
check(assert_type(i0, "pd.Interval[int]"), pd.Interval)
19+
check(assert_type(i1, "pd.Interval[float]"), pd.Interval)
20+
check(assert_type(i2, "pd.Interval[pd.Timestamp]"), pd.Interval)
21+
check(assert_type(i3, "pd.Interval[pd.Timedelta]"), pd.Interval)
22+
23+
check(assert_type(i0.closed, Literal["left", "right", "both", "neither"]), str)
24+
check(assert_type(i0.closed_left, bool), bool)
25+
check(assert_type(i0.closed_right, bool), bool)
26+
check(assert_type(i0.is_empty, bool), bool)
27+
check(assert_type(i0.left, int), int)
28+
check(assert_type(i0.length, int), int)
29+
check(assert_type(i0.mid, float), float)
30+
check(assert_type(i0.open_left, bool), bool)
31+
check(assert_type(i0.open_right, bool), bool)
32+
check(assert_type(i0.right, int), int)
33+
34+
check(assert_type(i1.closed, Literal["left", "right", "both", "neither"]), str)
35+
check(assert_type(i1.closed_left, bool), bool)
36+
check(assert_type(i1.closed_right, bool), bool)
37+
check(assert_type(i1.is_empty, bool), bool)
38+
check(assert_type(i1.left, float), float)
39+
check(assert_type(i1.length, float), float)
40+
check(assert_type(i1.mid, float), float)
41+
check(assert_type(i1.open_left, bool), bool)
42+
check(assert_type(i1.open_right, bool), bool)
43+
check(assert_type(i1.right, float), float)
44+
45+
check(assert_type(i2.closed, Literal["left", "right", "both", "neither"]), str)
46+
check(assert_type(i2.closed_left, bool), bool)
47+
check(assert_type(i2.closed_right, bool), bool)
48+
check(assert_type(i2.is_empty, bool), bool)
49+
check(assert_type(i2.left, pd.Timestamp), pd.Timestamp)
50+
check(assert_type(i2.length, pd.Timedelta), pd.Timedelta)
51+
check(assert_type(i2.mid, pd.Timestamp), pd.Timestamp)
52+
check(assert_type(i2.open_left, bool), bool)
53+
check(assert_type(i2.open_right, bool), bool)
54+
check(assert_type(i2.right, pd.Timestamp), pd.Timestamp)
55+
56+
check(assert_type(i3.closed, Literal["left", "right", "both", "neither"]), str)
57+
check(assert_type(i3.closed_left, bool), bool)
58+
check(assert_type(i3.closed_right, bool), bool)
59+
check(assert_type(i3.is_empty, bool), bool)
60+
check(assert_type(i3.left, pd.Timedelta), pd.Timedelta)
61+
check(assert_type(i3.length, pd.Timedelta), pd.Timedelta)
62+
check(assert_type(i3.mid, pd.Timedelta), pd.Timedelta)
63+
check(assert_type(i3.open_left, bool), bool)
64+
check(assert_type(i3.open_right, bool), bool)
65+
check(assert_type(i3.right, pd.Timedelta), pd.Timedelta)
66+
67+
check(assert_type(i0.overlaps(pd.Interval(0.5, 1.5, closed="left")), bool), bool)
68+
check(assert_type(i0.overlaps(pd.Interval(2, 3, closed="left")), bool), bool)
69+
70+
check(assert_type(i1.overlaps(pd.Interval(0.5, 1.5, closed="left")), bool), bool)
71+
check(assert_type(i1.overlaps(pd.Interval(2, 3, closed="left")), bool), bool)
72+
ts1 = pd.Timestamp(year=2017, month=1, day=1)
73+
ts2 = pd.Timestamp(year=2017, month=1, day=2)
74+
check(assert_type(i2.overlaps(pd.Interval(ts1, ts2, closed="left")), bool), bool)
75+
td1 = pd.Timedelta(days=1)
76+
td2 = pd.Timedelta(days=3)
77+
check(assert_type(i3.overlaps(pd.Interval(td1, td2, closed="left")), bool), bool)
78+
79+
check(assert_type(i0 * 3, "pd.Interval[int]"), pd.Interval)
80+
check(assert_type(i1 * 3, "pd.Interval[float]"), pd.Interval)
81+
check(assert_type(i3 * 3, "pd.Interval[pd.Timedelta]"), pd.Interval)
82+
83+
check(assert_type(i0 * 3.5, "pd.Interval[float]"), pd.Interval)
84+
check(assert_type(i1 * 3.5, "pd.Interval[float]"), pd.Interval)
85+
check(assert_type(i3 * 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval)
86+
87+
check(assert_type(3 * i0, "pd.Interval[int]"), pd.Interval)
88+
check(assert_type(3 * i1, "pd.Interval[float]"), pd.Interval)
89+
check(assert_type(3 * i3, "pd.Interval[pd.Timedelta]"), pd.Interval)
90+
91+
check(assert_type(3.5 * i0, "pd.Interval[float]"), pd.Interval)
92+
check(assert_type(3.5 * i1, "pd.Interval[float]"), pd.Interval)
93+
check(assert_type(3.5 * i3, "pd.Interval[pd.Timedelta]"), pd.Interval)
94+
95+
check(assert_type(i0 / 3, "pd.Interval[int]"), pd.Interval)
96+
check(assert_type(i1 / 3, "pd.Interval[float]"), pd.Interval)
97+
check(assert_type(i3 / 3, "pd.Interval[pd.Timedelta]"), pd.Interval)
98+
99+
check(assert_type(i0 / 3.5, "pd.Interval[float]"), pd.Interval)
100+
check(assert_type(i1 / 3.5, "pd.Interval[float]"), pd.Interval)
101+
check(assert_type(i3 / 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval)
102+
103+
check(assert_type(i0 // 3, "pd.Interval[int]"), pd.Interval)
104+
check(assert_type(i1 // 3, "pd.Interval[float]"), pd.Interval)
105+
check(assert_type(i3 // 3, "pd.Interval[pd.Timedelta]"), pd.Interval)
106+
107+
check(assert_type(i0 // 3.5, "pd.Interval[float]"), pd.Interval)
108+
check(assert_type(i1 // 3.5, "pd.Interval[float]"), pd.Interval)
109+
check(assert_type(i3 // 3.5, "pd.Interval[pd.Timedelta]"), pd.Interval)
110+
111+
check(assert_type(i0 - 1, "pd.Interval[int]"), pd.Interval)
112+
check(assert_type(i1 - 1, "pd.Interval[float]"), pd.Interval)
113+
check(
114+
assert_type(i2 - pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval
115+
)
116+
check(
117+
assert_type(i3 - pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval
118+
)
119+
120+
check(assert_type(i0 - 1.5, "pd.Interval[float]"), pd.Interval)
121+
check(assert_type(i1 - 1.5, "pd.Interval[float]"), pd.Interval)
122+
check(
123+
assert_type(i2 - pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval
124+
)
125+
check(
126+
assert_type(i3 - pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval
127+
)
128+
129+
check(assert_type(i0 + 1, "pd.Interval[int]"), pd.Interval)
130+
check(assert_type(i1 + 1, "pd.Interval[float]"), pd.Interval)
131+
check(
132+
assert_type(i2 + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval
133+
)
134+
check(
135+
assert_type(i3 + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval
136+
)
137+
138+
check(assert_type(i0 + 1.5, "pd.Interval[float]"), pd.Interval)
139+
check(assert_type(i1 + 1.5, "pd.Interval[float]"), pd.Interval)
140+
check(
141+
assert_type(i2 + pd.Timedelta(days=1), "pd.Interval[pd.Timestamp]"), pd.Interval
142+
)
143+
check(
144+
assert_type(i3 + pd.Timedelta(days=1), "pd.Interval[pd.Timedelta]"), pd.Interval
145+
)
146+
147+
check(assert_type(0.5 in i0, bool), bool)
148+
check(assert_type(1 in i0, bool), bool)
149+
check(assert_type(1 in i1, bool), bool)
150+
check(assert_type(pd.Timestamp("2000-1-1") in i2, bool), bool)
151+
check(assert_type(pd.Timedelta(days=1) in i3, bool), bool)
152+
153+
check(assert_type(hash(i0), int), int)
154+
check(assert_type(hash(i1), int), int)
155+
check(assert_type(hash(i2), int), int)
156+
check(assert_type(hash(i3), int), int)

0 commit comments

Comments
 (0)