Skip to content

Commit d40f7e9

Browse files
MarcoGorelliMarcoGorelli
committed
type DatetimeIndex.indexer_at_time and DatetimeIndex.indexer_between_time 
1 parent 01d597f commit d40f7e9

File tree

3 files changed

+127
-88
lines changed

3 files changed

+127
-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: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
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 = pd.date_range("2022-01-01", "2022-01-31", freq="D")
20+
x = pd.Timestamp("2022-01-17")
21+
idx = pd.Index(data, name="date")
22+
check(assert_type(data[x <= idx], pd.DatetimeIndex), pd.DatetimeIndex)
23+
check(assert_type(data[x < idx], pd.DatetimeIndex), pd.DatetimeIndex)
24+
check(assert_type(data[x >= idx], pd.DatetimeIndex), pd.DatetimeIndex)
25+
check(assert_type(data[x > idx], pd.DatetimeIndex), pd.DatetimeIndex)
26+
check(assert_type(data[idx < x], pd.DatetimeIndex), pd.DatetimeIndex)
27+
check(assert_type(data[idx >= x], pd.DatetimeIndex), pd.DatetimeIndex)
28+
check(assert_type(data[idx > x], pd.DatetimeIndex), pd.DatetimeIndex)
29+
check(assert_type(data[idx <= x], pd.DatetimeIndex), pd.DatetimeIndex)
30+
31+
dt_idx = pd.DatetimeIndex(data, name="date")
32+
check(assert_type(data[x <= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
33+
check(assert_type(data[x >= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
34+
check(assert_type(data[x < dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
35+
check(assert_type(data[x > dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)
36+
check(assert_type(data[dt_idx <= x], pd.DatetimeIndex), pd.DatetimeIndex)
37+
check(assert_type(data[dt_idx >= x], pd.DatetimeIndex), pd.DatetimeIndex)
38+
check(assert_type(data[dt_idx < x], pd.DatetimeIndex), pd.DatetimeIndex)
39+
check(assert_type(data[dt_idx > x], pd.DatetimeIndex), pd.DatetimeIndex)
40+
41+
ind = pd.Index([1, 2, 3])
42+
check(assert_type(ind <= 2, np_1darray[np.bool]), np_1darray[np.bool])
43+
check(assert_type(ind >= 2, np_1darray[np.bool]), np_1darray[np.bool])
44+
check(assert_type(ind < 2, np_1darray[np.bool]), np_1darray[np.bool])
45+
check(assert_type(ind > 2, np_1darray[np.bool]), np_1darray[np.bool])
46+
47+
48+
def test_datetime_index_constructor() -> None:
49+
check(assert_type(pd.DatetimeIndex(["2020"]), pd.DatetimeIndex), pd.DatetimeIndex)
50+
check(
51+
assert_type(pd.DatetimeIndex(["2020"], name="ts"), pd.DatetimeIndex),
52+
pd.DatetimeIndex,
53+
)
54+
check(
55+
assert_type(pd.DatetimeIndex(["2020"], freq="D"), pd.DatetimeIndex),
56+
pd.DatetimeIndex,
57+
)
58+
check(
59+
assert_type(pd.DatetimeIndex(["2020"], tz="Asia/Kathmandu"), pd.DatetimeIndex),
60+
pd.DatetimeIndex,
61+
)
62+
63+
# https://github.com/microsoft/python-type-stubs/issues/115
64+
df = pd.DataFrame({"A": [1, 2, 3], "B": [5, 6, 7]})
65+
66+
check(
67+
assert_type(
68+
pd.DatetimeIndex(data=df["A"], tz=None, ambiguous="NaT", copy=True),
69+
pd.DatetimeIndex,
70+
),
71+
pd.DatetimeIndex,
72+
)
73+
74+
75+
def test_intersection() -> None:
76+
# GH 744
77+
index = pd.DatetimeIndex(["2022-01-01"])
78+
check(assert_type(index.intersection(index), pd.DatetimeIndex), pd.DatetimeIndex)
79+
check(
80+
assert_type(index.intersection([pd.Timestamp("1/1/2023")]), pd.DatetimeIndex),
81+
pd.DatetimeIndex,
82+
)
83+
84+
85+
def test_datetime_index_max_min_reductions() -> None:
86+
dtidx = pd.DatetimeIndex(["2020-01-01", "2020-01-02"])
87+
check(assert_type(dtidx.argmax(), np.int64), np.int64)
88+
check(assert_type(dtidx.argmin(), np.int64), np.int64)
89+
check(assert_type(dtidx.max(), pd.Timestamp), pd.Timestamp)
90+
check(assert_type(dtidx.min(), pd.Timestamp), pd.Timestamp)
91+
92+
93+
def test_datetimeindex_shift() -> None:
94+
ind = pd.date_range("2023-01-01", "2023-02-01")
95+
check(assert_type(ind.shift(1), pd.DatetimeIndex), pd.DatetimeIndex)
96+
97+
98+
def test_datetimeindex_indexer_at_time() -> None:
99+
dti = pd.date_range("2023-01-01", "2023-02-01")
100+
check(assert_type(dti.indexer_at_time("10:00"), np_1darray[np.intp]), np_1darray)
101+
check(assert_type(dti.indexer_at_time(time(10)), np_1darray[np.intp]), np_1darray)
102+
103+
104+
def test_datetimeindex_indexer_between_time() -> None:
105+
dti = pd.date_range("2023-01-01", "2023-02-01")
106+
check(
107+
assert_type(
108+
dti.indexer_between_time(
109+
"10:00", time(11), include_start=False, include_end=True
110+
),
111+
np_1darray[np.intp],
112+
),
113+
np_1darray,
114+
)
115+
check(
116+
assert_type(dti.indexer_between_time(time(10), "11:00"), np_1darray[np.intp]),
117+
np_1darray,
118+
)

‎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)