diff --git a/doc/source/whatsnew/v1.2.0.rst b/doc/source/whatsnew/v1.2.0.rst index a3b5ba616b258..ff4335e921e3d 100644 --- a/doc/source/whatsnew/v1.2.0.rst +++ b/doc/source/whatsnew/v1.2.0.rst @@ -587,7 +587,7 @@ Interval - Bug in :meth:`DataFrame.replace` and :meth:`Series.replace` where :class:`Interval` dtypes would be converted to object dtypes (:issue:`34871`) - Bug in :meth:`IntervalIndex.take` with negative indices and ``fill_value=None`` (:issue:`37330`) -- +- Bug in :meth:`IntervalIndex.putmask` with datetime-like dtype incorrectly casting to object dtype (:issue:`37968`) - Indexing diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 07862e0b9bb48..b40531bd42af8 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -300,3 +300,24 @@ def __repr__(self) -> str: data = ",\n".join(lines) class_name = f"<{type(self).__name__}>" return f"{class_name}\n[\n{data}\n]\nShape: {self.shape}, dtype: {self.dtype}" + + # ------------------------------------------------------------------------ + # __array_function__ methods + + def putmask(self, mask, value): + """ + Analogue to np.putmask(self, mask, value) + + Parameters + ---------- + mask : np.ndarray[bool] + value : scalar or listlike + + Raises + ------ + TypeError + If value cannot be cast to self.dtype. + """ + value = self._validate_setitem_value(value) + + np.putmask(self._ndarray, mask, value) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 5209d83ade309..5eb890c9817c0 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -4344,11 +4344,9 @@ def putmask(self, mask, value): numpy.ndarray.putmask : Changes elements of an array based on conditional and input values. """ - values = self.values.copy() + values = self._values.copy() try: converted = self._validate_fill_value(value) - np.putmask(values, mask, converted) - return self._shallow_copy(values) except (ValueError, TypeError) as err: if is_object_dtype(self): raise err @@ -4356,6 +4354,9 @@ def putmask(self, mask, value): # coerces to object return self.astype(object).putmask(mask, value) + np.putmask(values, mask, converted) + return self._shallow_copy(values) + def equals(self, other: object) -> bool: """ Determine if two Index object are equal. diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py index c117c32f26d25..5db84a5d0a50a 100644 --- a/pandas/core/indexes/extension.py +++ b/pandas/core/indexes/extension.py @@ -359,15 +359,13 @@ def insert(self, loc: int, item): return type(self)._simple_new(new_arr, name=self.name) def putmask(self, mask, value): + res_values = self._data.copy() try: - value = self._data._validate_setitem_value(value) + res_values.putmask(mask, value) except (TypeError, ValueError): return self.astype(object).putmask(mask, value) - new_values = self._data._ndarray.copy() - np.putmask(new_values, mask, value) - new_arr = self._data._from_backing_data(new_values) - return type(self)._simple_new(new_arr, name=self.name) + return type(self)._simple_new(res_values, name=self.name) def _wrap_joined_index(self: _T, joined: np.ndarray, other: _T) -> _T: name = get_op_result_name(self, other) diff --git a/pandas/core/indexes/interval.py b/pandas/core/indexes/interval.py index b0f8be986fe5d..3c96685638ee8 100644 --- a/pandas/core/indexes/interval.py +++ b/pandas/core/indexes/interval.py @@ -858,6 +858,22 @@ def mid(self): def length(self): return Index(self._data.length, copy=False) + def putmask(self, mask, value): + arr = self._data.copy() + try: + value_left, value_right = arr._validate_setitem_value(value) + except (ValueError, TypeError): + return self.astype(object).putmask(mask, value) + + if isinstance(self._data._left, np.ndarray): + np.putmask(arr._left, mask, value_left) + np.putmask(arr._right, mask, value_right) + else: + # TODO: special case not needed with __array_function__ + arr._left.putmask(mask, value_left) + arr._right.putmask(mask, value_right) + return type(self)._simple_new(arr, name=self.name) + @Appender(Index.where.__doc__) def where(self, cond, other=None): if other is None: diff --git a/pandas/core/indexes/numeric.py b/pandas/core/indexes/numeric.py index 9dacdd6dea9ca..24aaf5885fe0e 100644 --- a/pandas/core/indexes/numeric.py +++ b/pandas/core/indexes/numeric.py @@ -122,7 +122,7 @@ def _validate_fill_value(self, value): # force conversion to object # so we don't lose the bools raise TypeError - if isinstance(value, str): + elif isinstance(value, str) or lib.is_complex(value): raise TypeError return value diff --git a/pandas/tests/indexes/interval/test_base.py b/pandas/tests/indexes/interval/test_base.py index c316655fbda8a..343c3d2e145f6 100644 --- a/pandas/tests/indexes/interval/test_base.py +++ b/pandas/tests/indexes/interval/test_base.py @@ -80,6 +80,30 @@ def test_where(self, closed, klass): result = idx.where(klass(cond)) tm.assert_index_equal(result, expected) + @pytest.mark.parametrize("tz", ["US/Pacific", None]) + def test_putmask_dt64(self, tz): + # GH#37968 + dti = date_range("2016-01-01", periods=9, tz=tz) + idx = IntervalIndex.from_breaks(dti) + mask = np.zeros(idx.shape, dtype=bool) + mask[0:3] = True + + result = idx.putmask(mask, idx[-1]) + expected = IntervalIndex([idx[-1]] * 3 + list(idx[3:])) + tm.assert_index_equal(result, expected) + + def test_putmask_td64(self): + # GH#37968 + dti = date_range("2016-01-01", periods=9) + tdi = dti - dti[0] + idx = IntervalIndex.from_breaks(tdi) + mask = np.zeros(idx.shape, dtype=bool) + mask[0:3] = True + + result = idx.putmask(mask, idx[-1]) + expected = IntervalIndex([idx[-1]] * 3 + list(idx[3:])) + tm.assert_index_equal(result, expected) + def test_getitem_2d_deprecated(self): # GH#30588 multi-dim indexing is deprecated, but raising is also acceptable idx = self.create_index()