implement ensure_localized in datetimelikeArrayMixin (#24378)

Author: jbrockmendel

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()
 

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):