diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index 6d8e41900ce2d..c01b04991e52b 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -1079,6 +1079,41 @@ def _evaluate_compare(self, other, op): result[mask] = filler return result + def _ensure_localized(self, arg, ambiguous='raise', nonexistent='raise', + from_utc=False): + """ + Ensure that we are re-localized. + + This is for compat as we can then call this on all datetimelike + arrays generally (ignored for Period/Timedelta) + + Parameters + ---------- + arg : Union[DatetimeLikeArray, DatetimeIndexOpsMixin, ndarray] + ambiguous : str, bool, or bool-ndarray, default 'raise' + nonexistent : str, default 'raise' + from_utc : bool, default False + If True, localize the i8 ndarray to UTC first before converting to + the appropriate tz. If False, localize directly to the tz. + + Returns + ------- + localized array + """ + + # reconvert to local tz + tz = getattr(self, 'tz', None) + if tz is not None: + if not isinstance(arg, type(self)): + arg = self._simple_new(arg) + if from_utc: + arg = arg.tz_localize('UTC').tz_convert(self.tz) + else: + arg = arg.tz_localize( + self.tz, ambiguous=ambiguous, nonexistent=nonexistent + ) + return arg + DatetimeLikeArrayMixin._add_comparison_ops() diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index dd2537c11a94c..db0cb88b06b2b 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -111,36 +111,17 @@ def _evaluate_compare(self, other, op): def _ensure_localized(self, arg, ambiguous='raise', nonexistent='raise', from_utc=False): - """ - Ensure that we are re-localized. - - This is for compat as we can then call this on all datetimelike - indexes generally (ignored for Period/Timedelta) - - Parameters - ---------- - arg : DatetimeIndex / i8 ndarray - ambiguous : str, bool, or bool-ndarray, default 'raise' - nonexistent : str, default 'raise' - from_utc : bool, default False - If True, localize the i8 ndarray to UTC first before converting to - the appropriate tz. If False, localize directly to the tz. - - Returns - ------- - localized DTI - """ - - # reconvert to local tz - if getattr(self, 'tz', None) is not None: - if not isinstance(arg, ABCIndexClass): - arg = self._simple_new(arg) - if from_utc: - arg = arg.tz_localize('UTC').tz_convert(self.tz) - else: - arg = arg.tz_localize( - self.tz, ambiguous=ambiguous, nonexistent=nonexistent - ) + # See DatetimeLikeArrayMixin._ensure_localized.__doc__ + + if getattr(self, 'tz', None): + # ensure_localized is only relevant for tz-aware DTI + from pandas.core.arrays import DatetimeArrayMixin as DatetimeArray + dtarr = DatetimeArray(self) + result = dtarr._ensure_localized(arg, + ambiguous=ambiguous, + nonexistent=nonexistent, + from_utc=from_utc) + return type(self)(result, name=self.name) return arg def _box_values_as_index(self):