Skip to content

Commit 95ce46c

Browse files
MarcoGorelliMarcoGorelli
authored
type DatetimeIndex.indexer_at_time and DatetimeIndex.indexer_between_time (#1438)
* type `DatetimeIndex.indexer_at_time` and `DatetimeIndex.indexer_between_time` * reorder, use `check/assert_type` when assigning to `data` and `idx`, add `snap` test * add dedupe comment
1 parent e1ad2eb commit 95ce46c

File tree

3 files changed

+153
-88
lines changed

3 files changed

+153
-88
lines changed

‎pandas-stubs/core/indexes/datetimes.pyi‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ from collections.abc import (
44
)
55
from datetime import (
66
datetime,
7+
time,
78
timedelta,
89
tzinfo as _tzinfo,
910
)
@@ -32,6 +33,7 @@ from pandas._typing import (
3233
IntervalClosedType,
3334
TimeUnit,
3435
TimeZones,
36+
np_1darray,
3537
np_ndarray_dt,
3638
np_ndarray_td,
3739
)
@@ -55,7 +57,6 @@ class DatetimeIndex(
5557
copy: bool = ...,
5658
name: Hashable = ...,
5759
) -> Self: ...
58-
def __reduce__(self): ...
5960

6061
# various ignores needed for mypy, as we do want to restrict what can be used in
6162
# arithmetic for these types
@@ -78,18 +79,19 @@ class DatetimeIndex(
7879
def to_series(
7980
self, index: Index | None = None, name: Hashable | None = None
8081
) -> Series[Timestamp]: ...
81-
def snap(self, freq: str = ...): ...
82-
def slice_indexer(self, start=..., end=..., step=...): ...
82+
def snap(self, freq: Frequency = "S") -> Self: ...
8383
@property
8484
def inferred_type(self) -> str: ...
85-
def indexer_at_time(self, time, asof: bool = ...): ...
85+
def indexer_at_time(
86+
self, time: str | time, asof: bool = False
87+
) -> np_1darray[np.intp]: ...
8688
def indexer_between_time(
8789
self,
88-
start_time: datetime | str,
89-
end_time: datetime | str,
90+
start_time: time | str,
91+
end_time: time | str,
9092
include_start: bool = True,
9193
include_end: bool = True,
92-
): ...
94+
) -> np_1darray[np.intp]: ...
9395
def to_julian_date(self) -> Index[float]: ...
9496
def isocalendar(self) -> DataFrame: ...
9597
@property

‎tests/indexes/test_datetime_index.py‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
from __future__ import annotations
2+
3+
from datetime import time
4+
5+
import numpy as np
6+
import pandas as pd
7+
from typing_extensions import (
8+
assert_type,
9+
)
10+
11+
from tests import (
12+
check,
13+
np_1darray,
14+
)
15+
16+
17+
def test_index_relops() -> None:
18+
# GH 265
19+
data = check(
20+
assert_type(
21+
pd.date_range("2022-01-01", "2022-01-31", freq="D"), pd.DatetimeIndex
22+
),
23+
pd.DatetimeIndex,
24+
)
25+
x = pd.Timestamp("2022-01-17")
26+
idx = check(
27+
assert_type(pd.Index(data, name="date"), "pd.Index[pd.Timestamp]"), pd.Index
28+
)
29+
check(assert_type(data[x <= idx], pd.DatetimeIndex), pd.DatetimeIndex)
30+
check(assert_type(data[x < idx], pd.DatetimeIndex), pd.DatetimeIndex)
31+
check(assert_type(data[x >= idx], pd.DatetimeIndex), pd.DatetimeIndex)
32+
check(assert_type(data[x > idx], pd.DatetimeIndex), pd.DatetimeIndex)
33+
check(assert_type(data[idx <= x], pd.DatetimeIndex), pd.DatetimeIndex)
34+
check(assert_type(data[idx < x], pd.DatetimeIndex), pd.DatetimeIndex)
35+
check(assert_type(data[idx >= x], pd.DatetimeIndex), pd.DatetimeIndex)
36+
check(assert_type(data[idx > x], pd.DatetimeIndex), pd.DatetimeIndex)
37+
38+
# TODO: https://github.com/pandas-dev/pandas-stubs/pull/1438#discussion_r2451864012
39+
# Can this be de-duplicated?
40+
dt_idx = pd.DatetimeIndex(data, name="date")
41+
check(assert_type(data[x <= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
42+
check(assert_type(data[x < dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
43+
check(assert_type(data[x >= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
44+
check(assert_type(data[x > dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
45+
check(assert_type(data[dt_idx <= x], pd.DatetimeIndex), pd.DatetimeIndex)
46+
check(assert_type(data[dt_idx < x], pd.DatetimeIndex), pd.DatetimeIndex)
47+
check(assert_type(data[dt_idx >= x], pd.DatetimeIndex), pd.DatetimeIndex)
48+
check(assert_type(data[dt_idx > x], pd.DatetimeIndex), pd.DatetimeIndex)
49+
50+
ind = pd.Index([1, 2, 3])
51+
check(assert_type(ind <= 2, np_1darray[np.bool]), np_1darray[np.bool])
52+
check(assert_type(ind < 2, np_1darray[np.bool]), np_1darray[np.bool])
53+
check(assert_type(ind >= 2, np_1darray[np.bool]), np_1darray[np.bool])
54+
check(assert_type(ind > 2, np_1darray[np.bool]), np_1darray[np.bool])
55+
56+
57+
def test_datetime_index_constructor() -> None:
58+
check(assert_type(pd.DatetimeIndex(["2020"]), pd.DatetimeIndex), pd.DatetimeIndex)
59+
check(
60+
assert_type(pd.DatetimeIndex(["2020"], name="ts"), pd.DatetimeIndex),
61+
pd.DatetimeIndex,
62+
)
63+
check(
64+
assert_type(pd.DatetimeIndex(["2020"], freq="D"), pd.DatetimeIndex),
65+
pd.DatetimeIndex,
66+
)
67+
check(
68+
assert_type(pd.DatetimeIndex(["2020"], tz="Asia/Kathmandu"), pd.DatetimeIndex),
69+
pd.DatetimeIndex,
70+
)
71+
72+
# https://github.com/microsoft/python-type-stubs/issues/115
73+
df = pd.DataFrame({"A": [1, 2, 3], "B": [5, 6, 7]})
74+
75+
check(
76+
assert_type(
77+
pd.DatetimeIndex(data=df["A"], tz=None, ambiguous="NaT", copy=True),
78+
pd.DatetimeIndex,
79+
),
80+
pd.DatetimeIndex,
81+
)
82+
83+
84+
def test_intersection() -> None:
85+
# GH 744
86+
index = pd.DatetimeIndex(["2022-01-01"])
87+
check(assert_type(index.intersection(index), pd.DatetimeIndex), pd.DatetimeIndex)
88+
check(
89+
assert_type(index.intersection([pd.Timestamp("1/1/2023")]), pd.DatetimeIndex),
90+
pd.DatetimeIndex,
91+
)
92+
93+
94+
def test_datetime_index_max_min_reductions() -> None:
95+
dtidx = pd.DatetimeIndex(["2020-01-01", "2020-01-02"])
96+
check(assert_type(dtidx.argmax(), np.int64), np.int64)
97+
check(assert_type(dtidx.argmin(), np.int64), np.int64)
98+
check(assert_type(dtidx.max(), pd.Timestamp), pd.Timestamp)
99+
check(assert_type(dtidx.min(), pd.Timestamp), pd.Timestamp)
100+
101+
102+
def test_datetimeindex_shift() -> None:
103+
ind = pd.date_range("2023-01-01", "2023-02-01")
104+
check(assert_type(ind.shift(1), pd.DatetimeIndex), pd.DatetimeIndex)
105+
106+
107+
def test_datetimeindex_indexer_at_time() -> None:
108+
dti = pd.date_range("2023-01-01", "2023-02-01")
109+
check(
110+
assert_type(dti.indexer_at_time("10:00"), np_1darray[np.intp]),
111+
np_1darray[np.intp],
112+
)
113+
check(
114+
assert_type(dti.indexer_at_time(time(10)), np_1darray[np.intp]),
115+
np_1darray[np.intp],
116+
)
117+
118+
119+
def test_datetimeindex_indexer_between_time() -> None:
120+
dti = pd.date_range("2023-01-01", "2023-02-01")
121+
check(
122+
assert_type(
123+
dti.indexer_between_time(
124+
"10:00", time(11), include_start=False, include_end=True
125+
),
126+
np_1darray[np.intp],
127+
),
128+
np_1darray[np.intp],
129+
)
130+
check(
131+
assert_type(dti.indexer_between_time(time(10), "11:00"), np_1darray[np.intp]),
132+
np_1darray[np.intp],
133+
)
134+
135+
136+
def test_datetimeindex_snap() -> None:
137+
dti = pd.date_range("2023-01-01", "2023-02-01")
138+
check(
139+
assert_type(
140+
dti.snap("MS"),
141+
pd.DatetimeIndex,
142+
),
143+
pd.DatetimeIndex,
144+
)

‎tests/indexes/test_indexes.py‎

Lines changed: 0 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -275,37 +275,6 @@ def test_index_arithmetic() -> None:
275275
check(assert_type(3 // idx, "pd.Index[float]"), pd.Index, np.float64)
276276

277277

278-
def test_index_relops() -> None:
279-
# GH 265
280-
data = pd.date_range("2022-01-01", "2022-01-31", freq="D")
281-
x = pd.Timestamp("2022-01-17")
282-
idx = pd.Index(data, name="date")
283-
check(assert_type(data[x <= idx], pd.DatetimeIndex), pd.DatetimeIndex)
284-
check(assert_type(data[x < idx], pd.DatetimeIndex), pd.DatetimeIndex)
285-
check(assert_type(data[x >= idx], pd.DatetimeIndex), pd.DatetimeIndex)
286-
check(assert_type(data[x > idx], pd.DatetimeIndex), pd.DatetimeIndex)
287-
check(assert_type(data[idx < x], pd.DatetimeIndex), pd.DatetimeIndex)
288-
check(assert_type(data[idx >= x], pd.DatetimeIndex), pd.DatetimeIndex)
289-
check(assert_type(data[idx > x], pd.DatetimeIndex), pd.DatetimeIndex)
290-
check(assert_type(data[idx <= x], pd.DatetimeIndex), pd.DatetimeIndex)
291-
292-
dt_idx = pd.DatetimeIndex(data, name="date")
293-
check(assert_type(data[x <= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
294-
check(assert_type(data[x >= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
295-
check(assert_type(data[x < dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
296-
check(assert_type(data[x > dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
297-
check(assert_type(data[dt_idx <= x], pd.DatetimeIndex), pd.DatetimeIndex)
298-
check(assert_type(data[dt_idx >= x], pd.DatetimeIndex), pd.DatetimeIndex)
299-
check(assert_type(data[dt_idx < x], pd.DatetimeIndex), pd.DatetimeIndex)
300-
check(assert_type(data[dt_idx > x], pd.DatetimeIndex), pd.DatetimeIndex)
301-
302-
ind = pd.Index([1, 2, 3])
303-
check(assert_type(ind <= 2, np_1darray[np.bool]), np_1darray[np.bool])
304-
check(assert_type(ind >= 2, np_1darray[np.bool]), np_1darray[np.bool])
305-
check(assert_type(ind < 2, np_1darray[np.bool]), np_1darray[np.bool])
306-
check(assert_type(ind > 2, np_1darray[np.bool]), np_1darray[np.bool])
307-
308-
309278
def test_range_index_union() -> None:
310279
check(
311280
assert_type(
@@ -1167,33 +1136,6 @@ def test_index_constructors() -> None:
11671136
pd.Index(flist, dtype=np.float16)
11681137

11691138

1170-
def test_datetime_index_constructor() -> None:
1171-
check(assert_type(pd.DatetimeIndex(["2020"]), pd.DatetimeIndex), pd.DatetimeIndex)
1172-
check(
1173-
assert_type(pd.DatetimeIndex(["2020"], name="ts"), pd.DatetimeIndex),
1174-
pd.DatetimeIndex,
1175-
)
1176-
check(
1177-
assert_type(pd.DatetimeIndex(["2020"], freq="D"), pd.DatetimeIndex),
1178-
pd.DatetimeIndex,
1179-
)
1180-
check(
1181-
assert_type(pd.DatetimeIndex(["2020"], tz="Asia/Kathmandu"), pd.DatetimeIndex),
1182-
pd.DatetimeIndex,
1183-
)
1184-
1185-
# https://github.com/microsoft/python-type-stubs/issues/115
1186-
df = pd.DataFrame({"A": [1, 2, 3], "B": [5, 6, 7]})
1187-
1188-
check(
1189-
assert_type(
1190-
pd.DatetimeIndex(data=df["A"], tz=None, ambiguous="NaT", copy=True),
1191-
pd.DatetimeIndex,
1192-
),
1193-
pd.DatetimeIndex,
1194-
)
1195-
1196-
11971139
def test_iter() -> None:
11981140
# GH 723
11991141
with pytest_warns_bounded(
@@ -1207,16 +1149,6 @@ def test_iter() -> None:
12071149
check(assert_type(ts, pd.Timestamp), pd.Timestamp)
12081150

12091151

1210-
def test_intersection() -> None:
1211-
# GH 744
1212-
index = pd.DatetimeIndex(["2022-01-01"])
1213-
check(assert_type(index.intersection(index), pd.DatetimeIndex), pd.DatetimeIndex)
1214-
check(
1215-
assert_type(index.intersection([pd.Timestamp("1/1/2023")]), pd.DatetimeIndex),
1216-
pd.DatetimeIndex,
1217-
)
1218-
1219-
12201152
def test_annotate() -> None:
12211153
# GH 502
12221154
df = pd.DataFrame({"a": [1, 2]})
@@ -1368,24 +1300,11 @@ def test_disallow_empty_index() -> None:
13681300
_0 = pd.Index() # type: ignore[call-overload] # pyright: ignore[reportCallIssue]
13691301

13701302

1371-
def test_datetime_index_max_min_reductions() -> None:
1372-
dtidx = pd.DatetimeIndex(["2020-01-01", "2020-01-02"])
1373-
check(assert_type(dtidx.argmax(), np.int64), np.int64)
1374-
check(assert_type(dtidx.argmin(), np.int64), np.int64)
1375-
check(assert_type(dtidx.max(), pd.Timestamp), pd.Timestamp)
1376-
check(assert_type(dtidx.min(), pd.Timestamp), pd.Timestamp)
1377-
1378-
13791303
def test_periodindex_shift() -> None:
13801304
ind = pd.period_range(start="2022-06-01", periods=10)
13811305
check(assert_type(ind.shift(1), pd.PeriodIndex), pd.PeriodIndex)
13821306

13831307

1384-
def test_datetimeindex_shift() -> None:
1385-
ind = pd.date_range("2023-01-01", "2023-02-01")
1386-
check(assert_type(ind.shift(1), pd.DatetimeIndex), pd.DatetimeIndex)
1387-
1388-
13891308
def test_timedeltaindex_shift() -> None:
13901309
ind = pd.date_range("1/1/2021", "1/5/2021") - pd.Timestamp("1/3/2019")
13911310
# broken on 3.0.0.dev0 as of 20250813, fix with pandas-dev/pandas/issues/62094

0 commit comments

Comments
 (0)