CLN: de-duplicate asof_locs (#37567)

Author: jbrockmendel

pandas/core/indexes/base.py

@@ -4503,7 +4503,7 @@ def asof(self, label):
                 loc = loc.indices(len(self))[-1]
             return self[loc]
 
-    def asof_locs(self, where, mask):
+    def asof_locs(self, where: "Index", mask) -> np.ndarray:
         """
         Return the locations (indices) of labels in the index.
 
@@ -4531,13 +4531,13 @@ def asof_locs(self, where, mask):
             which correspond to the return values of the `asof` function
             for every element in `where`.
         """
-        locs = self.values[mask].searchsorted(where.values, side="right")
+        locs = self._values[mask].searchsorted(where._values, side="right")
         locs = np.where(locs > 0, locs - 1, 0)
 
         result = np.arange(len(self))[mask].take(locs)
 
         first = mask.argmax()
-        result[(locs == 0) & (where.values < self.values[first])] = -1
+        result[(locs == 0) & (where._values < self._values[first])] = -1
 
         return result
 

pandas/core/indexes/period.py

@@ -404,28 +404,17 @@ def __array_wrap__(self, result, context=None):
         # cannot pass _simple_new as it is
         return type(self)(result, freq=self.freq, name=self.name)
 
-    def asof_locs(self, where, mask: np.ndarray) -> np.ndarray:
+    def asof_locs(self, where: Index, mask: np.ndarray) -> np.ndarray:
         """
         where : array of timestamps
         mask : array of booleans where data is not NA
         """
-        where_idx = where
-        if isinstance(where_idx, DatetimeIndex):
-            where_idx = PeriodIndex(where_idx._values, freq=self.freq)
-        elif not isinstance(where_idx, PeriodIndex):
+        if isinstance(where, DatetimeIndex):
+            where = PeriodIndex(where._values, freq=self.freq)
+        elif not isinstance(where, PeriodIndex):
             raise TypeError("asof_locs `where` must be DatetimeIndex or PeriodIndex")
-        elif where_idx.freq != self.freq:
-            raise raise_on_incompatible(self, where_idx)
 
-        locs = self.asi8[mask].searchsorted(where_idx.asi8, side="right")
-
-        locs = np.where(locs > 0, locs - 1, 0)
-        result = np.arange(len(self))[mask].take(locs)
-
-        first = mask.argmax()
-        result[(locs == 0) & (where_idx.asi8 < self.asi8[first])] = -1
-
-        return result
+        return super().asof_locs(where, mask)
 
     @doc(Index.astype)
     def astype(self, dtype, copy: bool = True, how="start"):