REF: simplify DTI/PI.get_loc (#49941)

* REF: simplify DTI.get_loc

* annotate

* update test message

Author: jbrockmendel

pandas/core/indexes/datetimes.py

@@ -526,29 +526,35 @@ def _parsed_string_to_bounds(self, reso: Resolution, parsed: dt.datetime):
                     "The index must be timezone aware when indexing "
                     "with a date string with a UTC offset"
                 )
-        start = self._maybe_cast_for_get_loc(start)
-        end = self._maybe_cast_for_get_loc(end)
+        # The flipped case with parsed.tz is None and self.tz is not None
+        #  is ruled out bc parsed and reso are produced by _parse_with_reso,
+        #  which localizes parsed.
         return start, end
 
-    def _disallow_mismatched_indexing(self, key, one_way: bool = False) -> None:
+    def _parse_with_reso(self, label: str):
+        parsed, reso = super()._parse_with_reso(label)
+
+        parsed = Timestamp(parsed)
+
+        if self.tz is not None and parsed.tzinfo is None:
+            # we special-case timezone-naive strings and timezone-aware
+            #  DatetimeIndex
+            # https://github.com/pandas-dev/pandas/pull/36148#issuecomment-687883081
+            parsed = parsed.tz_localize(self.tz)
+
+        return parsed, reso
+
+    def _disallow_mismatched_indexing(self, key) -> None:
         """
         Check for mismatched-tzawareness indexing and re-raise as KeyError.
         """
+        # we get here with isinstance(key, self._data._recognized_scalars)
         try:
-            self._deprecate_mismatched_indexing(key, one_way=one_way)
+            # GH#36148
+            self._data._assert_tzawareness_compat(key)
         except TypeError as err:
             raise KeyError(key) from err
 
-    def _deprecate_mismatched_indexing(self, key, one_way: bool = False) -> None:
-        # GH#36148
-        # we get here with isinstance(key, self._data._recognized_scalars)
-        if self.tz is not None and one_way:
-            # we special-case timezone-naive strings and timezone-aware
-            #  DatetimeIndex
-            return
-
-        self._data._assert_tzawareness_compat(key)
-
     def get_loc(self, key, method=None, tolerance=None):
         """
         Get integer location for requested label
@@ -566,15 +572,15 @@ def get_loc(self, key, method=None, tolerance=None):
         if isinstance(key, self._data._recognized_scalars):
             # needed to localize naive datetimes
             self._disallow_mismatched_indexing(key)
-            key = self._maybe_cast_for_get_loc(key)
+            key = Timestamp(key)
 
         elif isinstance(key, str):
 
             try:
                 parsed, reso = self._parse_with_reso(key)
             except ValueError as err:
                 raise KeyError(key) from err
-            self._disallow_mismatched_indexing(parsed, one_way=True)
+            self._disallow_mismatched_indexing(parsed)
 
             if self._can_partial_date_slice(reso):
                 try:
@@ -583,7 +589,7 @@ def get_loc(self, key, method=None, tolerance=None):
                     if method is None:
                         raise KeyError(key) from err
 
-            key = self._maybe_cast_for_get_loc(key)
+            key = parsed
 
         elif isinstance(key, dt.timedelta):
             # GH#20464
@@ -607,24 +613,6 @@ def get_loc(self, key, method=None, tolerance=None):
         except KeyError as err:
             raise KeyError(orig_key) from err
 
-    def _maybe_cast_for_get_loc(self, key) -> Timestamp:
-        # needed to localize naive datetimes or dates (GH 35690)
-        try:
-            key = Timestamp(key)
-        except ValueError as err:
-            # FIXME(dateutil#1180): we get here because parse_with_reso
-            #  doesn't raise on "t2m"
-            if not isinstance(key, str):
-                # Not expected to be reached, but check to be sure
-                raise  # pragma: no cover
-            raise KeyError(key) from err
-
-        if key.tzinfo is None:
-            key = key.tz_localize(self.tz)
-        else:
-            key = key.tz_convert(self.tz)
-        return key
-
     @doc(DatetimeTimedeltaMixin._maybe_cast_slice_bound)
     def _maybe_cast_slice_bound(self, label, side: str):
 
@@ -635,8 +623,8 @@ def _maybe_cast_slice_bound(self, label, side: str):
             label = Timestamp(label).to_pydatetime()
 
         label = super()._maybe_cast_slice_bound(label, side)
-        self._deprecate_mismatched_indexing(label)
-        return self._maybe_cast_for_get_loc(label)
+        self._data._assert_tzawareness_compat(label)
+        return Timestamp(label)
 
     def slice_indexer(self, start=None, end=None, step=None):
         """

pandas/core/indexes/period.py

@@ -420,18 +420,14 @@ def get_loc(self, key, method=None, tolerance=None):
             if reso == self._resolution_obj:
                 # the reso < self._resolution_obj case goes
                 #  through _get_string_slice
-                key = self._cast_partial_indexing_scalar(key)
-                loc = self.get_loc(key, method=method, tolerance=tolerance)
-                # Recursing instead of falling through matters for the exception
-                #  message in test_get_loc3 (though not clear if that really matters)
-                return loc
+                key = self._cast_partial_indexing_scalar(parsed)
             elif method is None:
                 raise KeyError(key)
             else:
                 key = self._cast_partial_indexing_scalar(parsed)
 
         elif isinstance(key, Period):
-            key = self._maybe_cast_for_get_loc(key)
+            self._disallow_mismatched_indexing(key)
 
         elif isinstance(key, datetime):
             key = self._cast_partial_indexing_scalar(key)
@@ -445,8 +441,7 @@ def get_loc(self, key, method=None, tolerance=None):
         except KeyError as err:
             raise KeyError(orig_key) from err
 
-    def _maybe_cast_for_get_loc(self, key: Period) -> Period:
-        # name is a misnomer, chosen for compat with DatetimeIndex
+    def _disallow_mismatched_indexing(self, key: Period) -> None:
         sfreq = self.freq
         kfreq = key.freq
         if not (
@@ -460,15 +455,14 @@ def _maybe_cast_for_get_loc(self, key: Period) -> Period:
             #  checking these two attributes is sufficient to check equality,
             #  and much more performant than `self.freq == key.freq`
             raise KeyError(key)
-        return key
 
-    def _cast_partial_indexing_scalar(self, label):
+    def _cast_partial_indexing_scalar(self, label: datetime) -> Period:
         try:
-            key = Period(label, freq=self.freq)
+            period = Period(label, freq=self.freq)
         except ValueError as err:
             # we cannot construct the Period
             raise KeyError(label) from err
-        return key
+        return period
 
     @doc(DatetimeIndexOpsMixin._maybe_cast_slice_bound)
     def _maybe_cast_slice_bound(self, label, side: str):

pandas/tests/indexes/period/test_indexing.py

@@ -373,7 +373,7 @@ def test_get_loc3(self):
         msg = "Input has different freq=None from PeriodArray\\(freq=D\\)"
         with pytest.raises(ValueError, match=msg):
             idx.get_loc("2000-01-10", method="nearest", tolerance="1 hour")
-        with pytest.raises(KeyError, match=r"^Period\('2000-01-10', 'D'\)$"):
+        with pytest.raises(KeyError, match=r"^'2000-01-10'$"):
             idx.get_loc("2000-01-10", method="nearest", tolerance="1 day")
         with pytest.raises(
             ValueError, match="list-like tolerance size must match target index size"