Added IntervalArray.__setitem__

Author: TomAugspurger

pandas/core/arrays/interval.py

@@ -456,6 +456,26 @@ def __getitem__(self, value):
 
         return self._shallow_copy(left, right)
 
+    def __setitem__(self, key, value):
+        if not (is_interval_dtype(value) or isinstance(value, Interval)):
+            msg = "'value' should be an interval type, got {} instead."
+            raise TypeError(msg.format(type(value)))
+
+        if value.closed != self.closed:
+            msg = "'value.closed' ({}) does not match {}."
+            raise ValueError(value.closed, self.closed)
+
+        # Need to ensure that left and right are updated atomically, so we're
+        # forced to copy, update the copy, and swap in the new values.
+        left = self.left.copy(deep=True)
+        right = self.right.copy(deep=True)
+
+        left.values[key] = value.left
+        right.values[key] = value.right
+
+        self._left = left
+        self._right = right
+
     def fillna(self, value=None, method=None, limit=None):
         if method is not None:
             raise TypeError('Filling by method is not supported for '

pandas/tests/extension/base/__init__.py

@@ -49,3 +49,4 @@ class TestMyDtype(BaseDtypeTests):
 from .methods import BaseMethodsTests  # noqa
 from .missing import BaseMissingTests  # noqa
 from .reshaping import BaseReshapingTests  # noqa
+from .setitem import BaseSetitemTests  # noqa

pandas/tests/extension/base/setitem.py

@@ -0,0 +1,38 @@
+import pytest
+
+from .base import BaseExtensionTests
+
+
+class BaseSetitemTests(BaseExtensionTests):
+    """Tests for ExtensionArray.__setitem__"""
+
+    def test_set_scalar(self, data):
+        expected = data.take([1, 1])
+        subset = data[:2].copy()
+
+        subset[0] = data[1]
+        self.assert_extension_array_equal(subset, expected)
+
+    def test_set_mask_scalar(self, data):
+        expected = data.take([1, 1, 2, 1])
+        subset = data[:4].copy()
+
+        subset[[True, True, False, True]] = data[1]
+        self.assert_extension_array_equal(subset, expected)
+
+    @pytest.mark.parametrize('key', [
+        [False, True, True, True],
+        [1, 2, 3],
+    ], ids=['mask', 'fancy'])
+    def test_set_array(self, key, data):
+        expected = data.take([0, 2, 2, 1])
+        value = data.take([2, 2, 1])
+        subset = data[:4].copy()
+
+        subset[key] = value
+        self.assert_extension_array_equal(subset, expected)
+
+    def test_bad_mask_bad_length_raise(self, data):
+        value = data[0]
+        with pytest.raises(IndexError):
+            data[[True, False]] = value

pandas/tests/extension/category/test_categorical.py

@@ -106,3 +106,7 @@ def test_value_counts(self, all_data, dropna):
 
 class TestCasting(base.BaseCastingTests):
     pass
+
+
+class TestSetitem(base.BaseSetitemTests):
+    pass

pandas/tests/extension/decimal/test_decimal.py

@@ -131,6 +131,10 @@ class TestGroupby(BaseDecimal, base.BaseGroupbyTests):
     pass
 
 
+class TestSetitem(BaseDecimal, base.BaseSetitemTests):
+    pass
+
+
 def test_series_constructor_coerce_data_to_extension_dtype_raises():
     xpr = ("Cannot cast data to extension dtype 'decimal'. Pass the "
            "extension array directly.")

pandas/tests/extension/test_interval.py

@@ -118,6 +118,10 @@ class TestReshaping(BaseInterval, base.BaseReshapingTests):
     pass
 
 
+class TestSetitem(BaseInterval, base.BaseSetitemTests):
+    pass
+
+
 def test_repr_matches():
     idx = IntervalIndex.from_breaks([1, 2, 3])
     a = repr(idx)