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"