shift and asof_locs too

Author: MarcoGorelli

pandas-stubs/core/indexes/period.pyi

@@ -52,14 +52,18 @@ class PeriodIndex(DatetimeIndexOpsMixin[pd.Period, np.object_], PeriodIndexField
     def __rsub__(  # pyright: ignore[reportIncompatibleMethodOverride]
         self, other: NaTType
     ) -> NaTType: ...
-    def asof_locs(self, where, mask): ...
+    def asof_locs(
+        self,
+        where: pd.DatetimeIndex | PeriodIndex,
+        mask: np.ndarray[tuple[int], np.dtype[np.bool_]],
+    ) -> Self: ...
     @property
     def is_full(self) -> bool: ...
     @property
     def inferred_type(self) -> str: ...
     @property
     def freqstr(self) -> str: ...
-    def shift(self, periods: int = 1, freq=...) -> Self: ...
+    def shift(self, periods: int = 1, freq: Frequency | None = None) -> Self: ...
 
 def period_range(
     start: (

tests/test_indexes.py

@@ -1502,3 +1502,13 @@ def test_period_index_constructor() -> None:
         ),
         pd.PeriodIndex,
     )
+
+
+def test_period_index_asof_locs() -> None:
+    idx = pd.PeriodIndex(["2000", "2001"], freq="D")
+    where = pd.DatetimeIndex(["2023-05-30 00:12:00", "2023-06-01 00:00:00"])
+    mask = np.ones(2, dtype=bool)
+    check(
+        assert_type(idx.asof_locs(where, mask), pd.PeriodIndex),
+        pd.PeriodIndex,
+    )