Skip to content

Commit fc92c6e

Browse files
authored
BUG: IntervalIndex.putmask with datetimelike dtypes (#37968)
1 parent 5c35871 commit fc92c6e

File tree

7 files changed

+70
-10
lines changed

7 files changed

+70
-10
lines changed

doc/source/whatsnew/v1.2.0.rst

+1-1
Original file line numberDiff line numberDiff line change
@@ -587,7 +587,7 @@ Interval
587587

588588
- Bug in :meth:`DataFrame.replace` and :meth:`Series.replace` where :class:`Interval` dtypes would be converted to object dtypes (:issue:`34871`)
589589
- Bug in :meth:`IntervalIndex.take` with negative indices and ``fill_value=None`` (:issue:`37330`)
590-
-
590+
- Bug in :meth:`IntervalIndex.putmask` with datetime-like dtype incorrectly casting to object dtype (:issue:`37968`)
591591
-
592592

593593
Indexing

pandas/core/arrays/_mixins.py

+21
Original file line numberDiff line numberDiff line change
@@ -300,3 +300,24 @@ def __repr__(self) -> str:
300300
data = ",\n".join(lines)
301301
class_name = f"<{type(self).__name__}>"
302302
return f"{class_name}\n[\n{data}\n]\nShape: {self.shape}, dtype: {self.dtype}"
303+
304+
# ------------------------------------------------------------------------
305+
# __array_function__ methods
306+
307+
def putmask(self, mask, value):
308+
"""
309+
Analogue to np.putmask(self, mask, value)
310+
311+
Parameters
312+
----------
313+
mask : np.ndarray[bool]
314+
value : scalar or listlike
315+
316+
Raises
317+
------
318+
TypeError
319+
If value cannot be cast to self.dtype.
320+
"""
321+
value = self._validate_setitem_value(value)
322+
323+
np.putmask(self._ndarray, mask, value)

pandas/core/indexes/base.py

+4-3
Original file line numberDiff line numberDiff line change
@@ -4344,18 +4344,19 @@ def putmask(self, mask, value):
43444344
numpy.ndarray.putmask : Changes elements of an array
43454345
based on conditional and input values.
43464346
"""
4347-
values = self.values.copy()
4347+
values = self._values.copy()
43484348
try:
43494349
converted = self._validate_fill_value(value)
4350-
np.putmask(values, mask, converted)
4351-
return self._shallow_copy(values)
43524350
except (ValueError, TypeError) as err:
43534351
if is_object_dtype(self):
43544352
raise err
43554353

43564354
# coerces to object
43574355
return self.astype(object).putmask(mask, value)
43584356

4357+
np.putmask(values, mask, converted)
4358+
return self._shallow_copy(values)
4359+
43594360
def equals(self, other: object) -> bool:
43604361
"""
43614362
Determine if two Index object are equal.

pandas/core/indexes/extension.py

+3-5
Original file line numberDiff line numberDiff line change
@@ -359,15 +359,13 @@ def insert(self, loc: int, item):
359359
return type(self)._simple_new(new_arr, name=self.name)
360360

361361
def putmask(self, mask, value):
362+
res_values = self._data.copy()
362363
try:
363-
value = self._data._validate_setitem_value(value)
364+
res_values.putmask(mask, value)
364365
except (TypeError, ValueError):
365366
return self.astype(object).putmask(mask, value)
366367

367-
new_values = self._data._ndarray.copy()
368-
np.putmask(new_values, mask, value)
369-
new_arr = self._data._from_backing_data(new_values)
370-
return type(self)._simple_new(new_arr, name=self.name)
368+
return type(self)._simple_new(res_values, name=self.name)
371369

372370
def _wrap_joined_index(self: _T, joined: np.ndarray, other: _T) -> _T:
373371
name = get_op_result_name(self, other)

pandas/core/indexes/interval.py

+16
Original file line numberDiff line numberDiff line change
@@ -858,6 +858,22 @@ def mid(self):
858858
def length(self):
859859
return Index(self._data.length, copy=False)
860860

861+
def putmask(self, mask, value):
862+
arr = self._data.copy()
863+
try:
864+
value_left, value_right = arr._validate_setitem_value(value)
865+
except (ValueError, TypeError):
866+
return self.astype(object).putmask(mask, value)
867+
868+
if isinstance(self._data._left, np.ndarray):
869+
np.putmask(arr._left, mask, value_left)
870+
np.putmask(arr._right, mask, value_right)
871+
else:
872+
# TODO: special case not needed with __array_function__
873+
arr._left.putmask(mask, value_left)
874+
arr._right.putmask(mask, value_right)
875+
return type(self)._simple_new(arr, name=self.name)
876+
861877
@Appender(Index.where.__doc__)
862878
def where(self, cond, other=None):
863879
if other is None:

pandas/core/indexes/numeric.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ def _validate_fill_value(self, value):
122122
# force conversion to object
123123
# so we don't lose the bools
124124
raise TypeError
125-
if isinstance(value, str):
125+
elif isinstance(value, str) or lib.is_complex(value):
126126
raise TypeError
127127

128128
return value

pandas/tests/indexes/interval/test_base.py

+24
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,30 @@ def test_where(self, closed, klass):
8080
result = idx.where(klass(cond))
8181
tm.assert_index_equal(result, expected)
8282

83+
@pytest.mark.parametrize("tz", ["US/Pacific", None])
84+
def test_putmask_dt64(self, tz):
85+
# GH#37968
86+
dti = date_range("2016-01-01", periods=9, tz=tz)
87+
idx = IntervalIndex.from_breaks(dti)
88+
mask = np.zeros(idx.shape, dtype=bool)
89+
mask[0:3] = True
90+
91+
result = idx.putmask(mask, idx[-1])
92+
expected = IntervalIndex([idx[-1]] * 3 + list(idx[3:]))
93+
tm.assert_index_equal(result, expected)
94+
95+
def test_putmask_td64(self):
96+
# GH#37968
97+
dti = date_range("2016-01-01", periods=9)
98+
tdi = dti - dti[0]
99+
idx = IntervalIndex.from_breaks(tdi)
100+
mask = np.zeros(idx.shape, dtype=bool)
101+
mask[0:3] = True
102+
103+
result = idx.putmask(mask, idx[-1])
104+
expected = IntervalIndex([idx[-1]] * 3 + list(idx[3:]))
105+
tm.assert_index_equal(result, expected)
106+
83107
def test_getitem_2d_deprecated(self):
84108
# GH#30588 multi-dim indexing is deprecated, but raising is also acceptable
85109
idx = self.create_index()

0 commit comments

Comments
 (0)