REF: share _partial_date_slice between PeriodIndex/DatetimeIndex (#31766)

Author: jbrockmendel

pandas/core/indexes/datetimelike.py

@@ -1,6 +1,7 @@
 """
 Base and utility classes for tseries type pandas objects.
 """
+from datetime import datetime
 from typing import Any, List, Optional, Union
 
 import numpy as np
@@ -411,6 +412,57 @@ def _convert_scalar_indexer(self, key, kind: str):
 
         return super()._convert_scalar_indexer(key, kind=kind)
 
+    def _validate_partial_date_slice(self, reso: str):
+        raise NotImplementedError
+
+    def _parsed_string_to_bounds(self, reso: str, parsed: datetime):
+        raise NotImplementedError
+
+    def _partial_date_slice(
+        self, reso: str, parsed: datetime, use_lhs: bool = True, use_rhs: bool = True
+    ):
+        """
+        Parameters
+        ----------
+        reso : str
+        parsed : datetime
+        use_lhs : bool, default True
+        use_rhs : bool, default True
+
+        Returns
+        -------
+        slice or ndarray[intp]
+        """
+        self._validate_partial_date_slice(reso)
+
+        t1, t2 = self._parsed_string_to_bounds(reso, parsed)
+        i8vals = self.asi8
+        unbox = self._data._unbox_scalar
+
+        if self.is_monotonic:
+
+            if len(self) and (
+                (use_lhs and t1 < self[0] and t2 < self[0])
+                or ((use_rhs and t1 > self[-1] and t2 > self[-1]))
+            ):
+                # we are out of range
+                raise KeyError
+
+            # TODO: does this depend on being monotonic _increasing_?
+
+            # a monotonic (sorted) series can be sliced
+            # Use asi8.searchsorted to avoid re-validating Periods/Timestamps
+            left = i8vals.searchsorted(unbox(t1), side="left") if use_lhs else None
+            right = i8vals.searchsorted(unbox(t2), side="right") if use_rhs else None
+            return slice(left, right)
+
+        else:
+            lhs_mask = (i8vals >= unbox(t1)) if use_lhs else True
+            rhs_mask = (i8vals <= unbox(t2)) if use_rhs else True
+
+            # try to find the dates
+            return (lhs_mask & rhs_mask).nonzero()[0]
+
     # --------------------------------------------------------------------
 
     __add__ = make_wrapped_arith_op("__add__")

pandas/core/indexes/datetimes.py

@@ -503,19 +503,9 @@ def _parsed_string_to_bounds(self, reso: str, parsed: datetime):
             end = end.tz_localize(self.tz)
         return start, end
 
-    def _partial_date_slice(
-        self, reso: str, parsed: datetime, use_lhs: bool = True, use_rhs: bool = True
-    ):
-        """
-        Parameters
-        ----------
-        reso : str
-        use_lhs : bool, default True
-        use_rhs : bool, default True
-        """
-        is_monotonic = self.is_monotonic
+    def _validate_partial_date_slice(self, reso: str):
         if (
-            is_monotonic
+            self.is_monotonic
             and reso in ["day", "hour", "minute", "second"]
             and self._resolution >= Resolution.get_reso(reso)
         ):
@@ -530,31 +520,6 @@ def _partial_date_slice(
             # _parsed_string_to_bounds allows it.
             raise KeyError
 
-        t1, t2 = self._parsed_string_to_bounds(reso, parsed)
-        stamps = self.asi8
-
-        if is_monotonic:
-
-            # we are out of range
-            if len(stamps) and (
-                (use_lhs and t1.value < stamps[0] and t2.value < stamps[0])
-                or ((use_rhs and t1.value > stamps[-1] and t2.value > stamps[-1]))
-            ):
-                raise KeyError
-
-            # a monotonic (sorted) series can be sliced
-            # Use asi8.searchsorted to avoid re-validating
-            left = stamps.searchsorted(t1.value, side="left") if use_lhs else None
-            right = stamps.searchsorted(t2.value, side="right") if use_rhs else None
-
-            return slice(left, right)
-
-        lhs_mask = (stamps >= t1.value) if use_lhs else True
-        rhs_mask = (stamps <= t2.value) if use_rhs else True
-
-        # try to find a the dates
-        return (lhs_mask & rhs_mask).nonzero()[0]
-
     def _maybe_promote(self, other):
         if other.inferred_type == "date":
             other = DatetimeIndex(other)

pandas/core/indexes/period.py

@@ -606,45 +606,24 @@ def _parsed_string_to_bounds(self, reso: str, parsed: datetime):
         iv = Period(parsed, freq=(grp, 1))
         return (iv.asfreq(self.freq, how="start"), iv.asfreq(self.freq, how="end"))
 
-    def _get_string_slice(self, key: str, use_lhs: bool = True, use_rhs: bool = True):
-        # TODO: Check for non-True use_lhs/use_rhs
-        parsed, reso = parse_time_string(key, self.freq)
+    def _validate_partial_date_slice(self, reso: str):
         grp = resolution.Resolution.get_freq_group(reso)
         freqn = resolution.get_freq_group(self.freq)
 
         if not grp < freqn:
             # TODO: we used to also check for
             #  reso in ["day", "hour", "minute", "second"]
             #  why is that check not needed?
-            raise ValueError(key)
-
-        t1, t2 = self._parsed_string_to_bounds(reso, parsed)
-        i8vals = self.asi8
-
-        if self.is_monotonic:
-
-            # we are out of range
-            if len(self) and (
-                (use_lhs and t1 < self[0] and t2 < self[0])
-                or ((use_rhs and t1 > self[-1] and t2 > self[-1]))
-            ):
-                raise KeyError(key)
-
-            # TODO: does this depend on being monotonic _increasing_?
-            #  If so, DTI will also be affected.
+            raise ValueError
 
-            # a monotonic (sorted) series can be sliced
-            # Use asi8.searchsorted to avoid re-validating Periods
-            left = i8vals.searchsorted(t1.ordinal, side="left") if use_lhs else None
-            right = i8vals.searchsorted(t2.ordinal, side="right") if use_rhs else None
-            return slice(left, right)
-
-        else:
-            lhs_mask = (i8vals >= t1.ordinal) if use_lhs else True
-            rhs_mask = (i8vals <= t2.ordinal) if use_rhs else True
+    def _get_string_slice(self, key: str, use_lhs: bool = True, use_rhs: bool = True):
+        # TODO: Check for non-True use_lhs/use_rhs
+        parsed, reso = parse_time_string(key, self.freq)
 
-            # try to find a the dates
-            return (lhs_mask & rhs_mask).nonzero()[0]
+        try:
+            return self._partial_date_slice(reso, parsed, use_lhs, use_rhs)
+        except KeyError:
+            raise KeyError(key)
 
     def _convert_tolerance(self, tolerance, target):
         tolerance = DatetimeIndexOpsMixin._convert_tolerance(self, tolerance, target)