reorder, use check/assert_type when assigning to data and idx, add

`snap` test

Author: MarcoGorelli

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