REF: share Index.putmask with ExtensionIndex (pandas-dev#43221)

Author: jbrockmendel

pandas/core/indexes/base.py

@@ -4857,6 +4857,7 @@ def _concat(self, to_concat: list[Index], name: Hashable) -> Index:
         result = concat_compat(to_concat_vals)
         return Index._with_infer(result, name=name)
 
+    @final
     def putmask(self, mask, value) -> Index:
         """
         Return a new Index of the values set with the mask.
@@ -4879,19 +4880,24 @@ def putmask(self, mask, value) -> Index:
         try:
             converted = self._validate_fill_value(value)
         except (ValueError, TypeError) as err:
-            if is_object_dtype(self):
+            if is_object_dtype(self):  # pragma: no cover
                 raise err
 
             dtype = self._find_common_type_compat(value)
             return self.astype(dtype).putmask(mask, value)
 
         values = self._values.copy()
-        # error: Argument 1 to "setitem_datetimelike_compat" has incompatible type
-        # "Union[ExtensionArray, ndarray]"; expected "ndarray"
-        converted = setitem_datetimelike_compat(
-            values, mask.sum(), converted  # type: ignore[arg-type]
-        )
-        np.putmask(values, mask, converted)
+
+        if isinstance(values, np.ndarray):
+            converted = setitem_datetimelike_compat(values, mask.sum(), converted)
+            np.putmask(values, mask, converted)
+
+        else:
+            # Note: we use the original value here, not converted, as
+            #  _validate_fill_value is not idempotent
+            # error: "ExtensionArray" has no attribute "putmask"
+            values.putmask(mask, value)  # type: ignore[attr-defined]
+
         return self._shallow_copy(values)
 
     def equals(self, other: Any) -> bool:

pandas/core/indexes/extension.py

@@ -31,7 +31,6 @@
     ABCSeries,
 )
 
-from pandas.core.array_algos.putmask import validate_putmask
 from pandas.core.arrays import (
     Categorical,
     DatetimeArray,
@@ -323,21 +322,6 @@ def searchsorted(self, value, side="left", sorter=None) -> np.ndarray:
         # overriding IndexOpsMixin improves performance GH#38083
         return self._data.searchsorted(value, side=side, sorter=sorter)
 
-    def putmask(self, mask, value) -> Index:
-        mask, noop = validate_putmask(self._data, mask)
-        if noop:
-            return self.copy()
-
-        try:
-            self._validate_fill_value(value)
-        except (ValueError, TypeError):
-            dtype = self._find_common_type_compat(value)
-            return self.astype(dtype).putmask(mask, value)
-
-        arr = self._data.copy()
-        arr.putmask(mask, value)
-        return type(self)._simple_new(arr, name=self.name)
-
     # ---------------------------------------------------------------------
 
     def _get_engine_target(self) -> np.ndarray: