From d40f7e952a2884372cdb08486f3af88a6cf43c45 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 22 Oct 2025 11:51:19 +0100 Subject: [PATCH 1/3] type `DatetimeIndex.indexer_at_time` and `DatetimeIndex.indexer_between_time` --- pandas-stubs/core/indexes/datetimes.pyi | 16 ++-- tests/indexes/test_datetime_index.py | 118 ++++++++++++++++++++++++ tests/indexes/test_indexes.py | 81 ---------------- 3 files changed, 127 insertions(+), 88 deletions(-) create mode 100644 tests/indexes/test_datetime_index.py diff --git a/pandas-stubs/core/indexes/datetimes.pyi b/pandas-stubs/core/indexes/datetimes.pyi index 4d30a8460..4ff9bc44f 100644 --- a/pandas-stubs/core/indexes/datetimes.pyi +++ b/pandas-stubs/core/indexes/datetimes.pyi @@ -4,6 +4,7 @@ from collections.abc import ( ) from datetime import ( datetime, + time, timedelta, tzinfo as _tzinfo, ) @@ -32,6 +33,7 @@ from pandas._typing import ( IntervalClosedType, TimeUnit, TimeZones, + np_1darray, np_ndarray_dt, np_ndarray_td, ) @@ -55,7 +57,6 @@ class DatetimeIndex( copy: bool = ..., name: Hashable = ..., ) -> Self: ... - def __reduce__(self): ... # various ignores needed for mypy, as we do want to restrict what can be used in # arithmetic for these types @@ -78,18 +79,19 @@ class DatetimeIndex( def to_series( self, index: Index | None = None, name: Hashable | None = None ) -> Series[Timestamp]: ... - def snap(self, freq: str = ...): ... - def slice_indexer(self, start=..., end=..., step=...): ... + def snap(self, freq: Frequency = "S") -> Self: ... @property def inferred_type(self) -> str: ... - def indexer_at_time(self, time, asof: bool = ...): ... + def indexer_at_time( + self, time: str | time, asof: bool = False + ) -> np_1darray[np.intp]: ... def indexer_between_time( self, - start_time: datetime | str, - end_time: datetime | str, + start_time: time | str, + end_time: time | str, include_start: bool = True, include_end: bool = True, - ): ... + ) -> np_1darray[np.intp]: ... def to_julian_date(self) -> Index[float]: ... def isocalendar(self) -> DataFrame: ... @property diff --git a/tests/indexes/test_datetime_index.py b/tests/indexes/test_datetime_index.py new file mode 100644 index 000000000..ce240561c --- /dev/null +++ b/tests/indexes/test_datetime_index.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from datetime import time + +import numpy as np +import pandas as pd +from typing_extensions import ( + assert_type, +) + +from tests import ( + check, + np_1darray, +) + + +def test_index_relops() -> None: + # GH 265 + data = pd.date_range("2022-01-01", "2022-01-31", freq="D") + x = pd.Timestamp("2022-01-17") + idx = pd.Index(data, name="date") + check(assert_type(data[x <= idx], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[x < idx], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[x >= idx], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[x > idx], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[idx < x], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[idx >= x], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[idx > x], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[idx <= x], pd.DatetimeIndex), pd.DatetimeIndex) + + dt_idx = pd.DatetimeIndex(data, name="date") + check(assert_type(data[x <= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[x >= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[x < dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[x > dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[dt_idx <= x], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[dt_idx >= x], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[dt_idx < x], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[dt_idx > x], pd.DatetimeIndex), pd.DatetimeIndex) + + ind = pd.Index([1, 2, 3]) + check(assert_type(ind <= 2, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(ind >= 2, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(ind < 2, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(ind > 2, np_1darray[np.bool]), np_1darray[np.bool]) + + +def test_datetime_index_constructor() -> None: + check(assert_type(pd.DatetimeIndex(["2020"]), pd.DatetimeIndex), pd.DatetimeIndex) + check( + assert_type(pd.DatetimeIndex(["2020"], name="ts"), pd.DatetimeIndex), + pd.DatetimeIndex, + ) + check( + assert_type(pd.DatetimeIndex(["2020"], freq="D"), pd.DatetimeIndex), + pd.DatetimeIndex, + ) + check( + assert_type(pd.DatetimeIndex(["2020"], tz="Asia/Kathmandu"), pd.DatetimeIndex), + pd.DatetimeIndex, + ) + + # https://github.com/microsoft/python-type-stubs/issues/115 + df = pd.DataFrame({"A": [1, 2, 3], "B": [5, 6, 7]}) + + check( + assert_type( + pd.DatetimeIndex(data=df["A"], tz=None, ambiguous="NaT", copy=True), + pd.DatetimeIndex, + ), + pd.DatetimeIndex, + ) + + +def test_intersection() -> None: + # GH 744 + index = pd.DatetimeIndex(["2022-01-01"]) + check(assert_type(index.intersection(index), pd.DatetimeIndex), pd.DatetimeIndex) + check( + assert_type(index.intersection([pd.Timestamp("1/1/2023")]), pd.DatetimeIndex), + pd.DatetimeIndex, + ) + + +def test_datetime_index_max_min_reductions() -> None: + dtidx = pd.DatetimeIndex(["2020-01-01", "2020-01-02"]) + check(assert_type(dtidx.argmax(), np.int64), np.int64) + check(assert_type(dtidx.argmin(), np.int64), np.int64) + check(assert_type(dtidx.max(), pd.Timestamp), pd.Timestamp) + check(assert_type(dtidx.min(), pd.Timestamp), pd.Timestamp) + + +def test_datetimeindex_shift() -> None: + ind = pd.date_range("2023-01-01", "2023-02-01") + check(assert_type(ind.shift(1), pd.DatetimeIndex), pd.DatetimeIndex) + + +def test_datetimeindex_indexer_at_time() -> None: + dti = pd.date_range("2023-01-01", "2023-02-01") + check(assert_type(dti.indexer_at_time("10:00"), np_1darray[np.intp]), np_1darray) + check(assert_type(dti.indexer_at_time(time(10)), np_1darray[np.intp]), np_1darray) + + +def test_datetimeindex_indexer_between_time() -> None: + dti = pd.date_range("2023-01-01", "2023-02-01") + check( + assert_type( + dti.indexer_between_time( + "10:00", time(11), include_start=False, include_end=True + ), + np_1darray[np.intp], + ), + np_1darray, + ) + check( + assert_type(dti.indexer_between_time(time(10), "11:00"), np_1darray[np.intp]), + np_1darray, + ) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 3b2960b3c..3ffe2d1cb 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -275,37 +275,6 @@ def test_index_arithmetic() -> None: check(assert_type(3 // idx, "pd.Index[float]"), pd.Index, np.float64) -def test_index_relops() -> None: - # GH 265 - data = pd.date_range("2022-01-01", "2022-01-31", freq="D") - x = pd.Timestamp("2022-01-17") - idx = pd.Index(data, name="date") - check(assert_type(data[x <= idx], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[x < idx], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[x >= idx], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[x > idx], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[idx < x], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[idx >= x], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[idx > x], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[idx <= x], pd.DatetimeIndex), pd.DatetimeIndex) - - dt_idx = pd.DatetimeIndex(data, name="date") - check(assert_type(data[x <= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[x >= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[x < dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[x > dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[dt_idx <= x], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[dt_idx >= x], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[dt_idx < x], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[dt_idx > x], pd.DatetimeIndex), pd.DatetimeIndex) - - ind = pd.Index([1, 2, 3]) - check(assert_type(ind <= 2, np_1darray[np.bool]), np_1darray[np.bool]) - check(assert_type(ind >= 2, np_1darray[np.bool]), np_1darray[np.bool]) - check(assert_type(ind < 2, np_1darray[np.bool]), np_1darray[np.bool]) - check(assert_type(ind > 2, np_1darray[np.bool]), np_1darray[np.bool]) - - def test_range_index_union() -> None: check( assert_type( @@ -1167,33 +1136,6 @@ def test_index_constructors() -> None: pd.Index(flist, dtype=np.float16) -def test_datetime_index_constructor() -> None: - check(assert_type(pd.DatetimeIndex(["2020"]), pd.DatetimeIndex), pd.DatetimeIndex) - check( - assert_type(pd.DatetimeIndex(["2020"], name="ts"), pd.DatetimeIndex), - pd.DatetimeIndex, - ) - check( - assert_type(pd.DatetimeIndex(["2020"], freq="D"), pd.DatetimeIndex), - pd.DatetimeIndex, - ) - check( - assert_type(pd.DatetimeIndex(["2020"], tz="Asia/Kathmandu"), pd.DatetimeIndex), - pd.DatetimeIndex, - ) - - # https://github.com/microsoft/python-type-stubs/issues/115 - df = pd.DataFrame({"A": [1, 2, 3], "B": [5, 6, 7]}) - - check( - assert_type( - pd.DatetimeIndex(data=df["A"], tz=None, ambiguous="NaT", copy=True), - pd.DatetimeIndex, - ), - pd.DatetimeIndex, - ) - - def test_iter() -> None: # GH 723 with pytest_warns_bounded( @@ -1207,16 +1149,6 @@ def test_iter() -> None: check(assert_type(ts, pd.Timestamp), pd.Timestamp) -def test_intersection() -> None: - # GH 744 - index = pd.DatetimeIndex(["2022-01-01"]) - check(assert_type(index.intersection(index), pd.DatetimeIndex), pd.DatetimeIndex) - check( - assert_type(index.intersection([pd.Timestamp("1/1/2023")]), pd.DatetimeIndex), - pd.DatetimeIndex, - ) - - def test_annotate() -> None: # GH 502 df = pd.DataFrame({"a": [1, 2]}) @@ -1368,24 +1300,11 @@ def test_disallow_empty_index() -> None: _0 = pd.Index() # type: ignore[call-overload] # pyright: ignore[reportCallIssue] -def test_datetime_index_max_min_reductions() -> None: - dtidx = pd.DatetimeIndex(["2020-01-01", "2020-01-02"]) - check(assert_type(dtidx.argmax(), np.int64), np.int64) - check(assert_type(dtidx.argmin(), np.int64), np.int64) - check(assert_type(dtidx.max(), pd.Timestamp), pd.Timestamp) - check(assert_type(dtidx.min(), pd.Timestamp), pd.Timestamp) - - def test_periodindex_shift() -> None: ind = pd.period_range(start="2022-06-01", periods=10) check(assert_type(ind.shift(1), pd.PeriodIndex), pd.PeriodIndex) -def test_datetimeindex_shift() -> None: - ind = pd.date_range("2023-01-01", "2023-02-01") - check(assert_type(ind.shift(1), pd.DatetimeIndex), pd.DatetimeIndex) - - def test_timedeltaindex_shift() -> None: ind = pd.date_range("1/1/2021", "1/5/2021") - pd.Timestamp("1/3/2019") # broken on 3.0.0.dev0 as of 20250813, fix with pandas-dev/pandas/issues/62094 From 35683fec1195ef95fe56597e2cadc8e9e6e1941a Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 22 Oct 2025 14:06:47 +0100 Subject: [PATCH 2/3] reorder, use `check/assert_type` when assigning to `data` and `idx`, add `snap` test --- tests/indexes/test_datetime_index.py | 44 +++++++++++++++++++++------- 1 file changed, 34 insertions(+), 10 deletions(-) diff --git a/tests/indexes/test_datetime_index.py b/tests/indexes/test_datetime_index.py index ce240561c..5ab0838da 100644 --- a/tests/indexes/test_datetime_index.py +++ b/tests/indexes/test_datetime_index.py @@ -16,32 +16,39 @@ def test_index_relops() -> None: # GH 265 - data = pd.date_range("2022-01-01", "2022-01-31", freq="D") + data = check( + assert_type( + pd.date_range("2022-01-01", "2022-01-31", freq="D"), pd.DatetimeIndex + ), + pd.DatetimeIndex, + ) x = pd.Timestamp("2022-01-17") - idx = pd.Index(data, name="date") + idx = check( + assert_type(pd.Index(data, name="date"), "pd.Index[pd.Timestamp]"), pd.Index + ) check(assert_type(data[x <= idx], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[x < idx], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[x >= idx], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[x > idx], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[idx <= x], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[idx < x], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[idx >= x], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[idx > x], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[idx <= x], pd.DatetimeIndex), pd.DatetimeIndex) dt_idx = pd.DatetimeIndex(data, name="date") check(assert_type(data[x <= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[x >= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[x < dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[x >= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[x > dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[dt_idx <= x], pd.DatetimeIndex), pd.DatetimeIndex) - check(assert_type(data[dt_idx >= x], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[dt_idx < x], pd.DatetimeIndex), pd.DatetimeIndex) + check(assert_type(data[dt_idx >= x], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[dt_idx > x], pd.DatetimeIndex), pd.DatetimeIndex) ind = pd.Index([1, 2, 3]) check(assert_type(ind <= 2, np_1darray[np.bool]), np_1darray[np.bool]) - check(assert_type(ind >= 2, np_1darray[np.bool]), np_1darray[np.bool]) check(assert_type(ind < 2, np_1darray[np.bool]), np_1darray[np.bool]) + check(assert_type(ind >= 2, np_1darray[np.bool]), np_1darray[np.bool]) check(assert_type(ind > 2, np_1darray[np.bool]), np_1darray[np.bool]) @@ -97,8 +104,14 @@ def test_datetimeindex_shift() -> None: def test_datetimeindex_indexer_at_time() -> None: dti = pd.date_range("2023-01-01", "2023-02-01") - check(assert_type(dti.indexer_at_time("10:00"), np_1darray[np.intp]), np_1darray) - check(assert_type(dti.indexer_at_time(time(10)), np_1darray[np.intp]), np_1darray) + check( + assert_type(dti.indexer_at_time("10:00"), np_1darray[np.intp]), + np_1darray[np.intp], + ) + check( + assert_type(dti.indexer_at_time(time(10)), np_1darray[np.intp]), + np_1darray[np.intp], + ) def test_datetimeindex_indexer_between_time() -> None: @@ -110,9 +123,20 @@ def test_datetimeindex_indexer_between_time() -> None: ), np_1darray[np.intp], ), - np_1darray, + np_1darray[np.intp], ) check( assert_type(dti.indexer_between_time(time(10), "11:00"), np_1darray[np.intp]), - np_1darray, + np_1darray[np.intp], + ) + + +def test_datetimeindex_snap() -> None: + dti = pd.date_range("2023-01-01", "2023-02-01") + check( + assert_type( + dti.snap("MS"), + pd.DatetimeIndex, + ), + pd.DatetimeIndex, ) From 94af7c313c2a386dcfe8de0b376866878b18a6df Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 23 Oct 2025 11:23:49 +0100 Subject: [PATCH 3/3] add dedupe comment --- tests/indexes/test_datetime_index.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/indexes/test_datetime_index.py b/tests/indexes/test_datetime_index.py index 5ab0838da..221a12607 100644 --- a/tests/indexes/test_datetime_index.py +++ b/tests/indexes/test_datetime_index.py @@ -35,6 +35,8 @@ def test_index_relops() -> None: check(assert_type(data[idx >= x], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[idx > x], pd.DatetimeIndex), pd.DatetimeIndex) + # TODO: https://github.com/pandas-dev/pandas-stubs/pull/1438#discussion_r2451864012 + # Can this be de-duplicated? dt_idx = pd.DatetimeIndex(data, name="date") check(assert_type(data[x <= dt_idx], pd.DatetimeIndex), pd.DatetimeIndex) check(assert_type(data[x < dt_idx], pd.DatetimeIndex), pd.DatetimeIndex)