BUG: Array.__setitem__ failing with nullable boolean mask

doc/source/whatsnew/v1.0.1.rst

@@ -70,6 +70,7 @@ Indexing
 
 -
 -
+- Bug where assigning to a ``Series`` using a IntegerArray / BooleanArray as a mask would raise ``TypeError`` (:issue:`31446`)
 
 Missing
 ^^^^^^^

pandas/core/arrays/boolean.py

@@ -26,6 +26,7 @@
 from pandas.core.dtypes.missing import isna, notna
 
 from pandas.core import nanops, ops
+from pandas.core.indexers import check_array_indexer
 
 from .masked import BaseMaskedArray
 
@@ -369,6 +370,7 @@ def __setitem__(self, key, value):
             value = value[0]
             mask = mask[0]
 
+        key = check_array_indexer(self, key)
         self._data[key] = value
         self._mask[key] = mask
 

pandas/core/arrays/integer.py

@@ -25,6 +25,7 @@
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import nanops, ops
+from pandas.core.indexers import check_array_indexer
 from pandas.core.ops import invalid_comparison
 from pandas.core.ops.common import unpack_zerodim_and_defer
 from pandas.core.tools.numeric import to_numeric
@@ -414,6 +415,7 @@ def __setitem__(self, key, value):
             value = value[0]
             mask = mask[0]
 
+        key = check_array_indexer(self, key)
         self._data[key] = value
         self._mask[key] = mask
 

pandas/core/arrays/string_.py

@@ -15,6 +15,7 @@
 from pandas.core import ops
 from pandas.core.arrays import PandasArray
 from pandas.core.construction import extract_array
+from pandas.core.indexers import check_array_indexer
 from pandas.core.missing import isna
 
 
@@ -224,6 +225,7 @@ def __setitem__(self, key, value):
             # extract_array doesn't extract PandasArray subclasses
             value = value._ndarray
 
+        key = check_array_indexer(self, key)
         scalar_key = lib.is_scalar(key)
         scalar_value = lib.is_scalar(value)
         if scalar_key and not scalar_value:

pandas/tests/arrays/test_integer.py

@@ -1074,6 +1074,23 @@ def test_cut(bins, right, include_lowest):
     tm.assert_categorical_equal(result, expected)
 
 
+def test_array_setitem_nullable_boolean_mask():
+    # GH 31446
+    ser = pd.Series([1, 2], dtype="Int64")
+    result = ser.where(ser > 1)
+    expected = pd.Series([pd.NA, 2], dtype="Int64")
+    tm.assert_series_equal(result, expected)
+
+
+def test_array_setitem():
+    # GH 31446
+    arr = pd.Series([1, 2], dtype="Int64").array
+    arr[arr > 1] = 1
+
+    expected = pd.array([1, 1], dtype="Int64")
+    tm.assert_extension_array_equal(arr, expected)
+
+
 # TODO(jreback) - these need testing / are broken
 
 # shift

pandas/tests/extension/base/setitem.py

@@ -195,3 +195,27 @@ def test_setitem_preserves_views(self, data):
         data[0] = data[1]
         assert view1[0] == data[1]
         assert view2[0] == data[1]
+
+    def test_setitem_nullable_integer(self, data):
+        # GH 31446
+        ser = pd.Series([1] * len(data) + [2] * len(data), dtype="Int64")
+        ser[ser > 1] = 3
+        expected = pd.Series([1] * len(data) + [3] * len(data), dtype="Int64")
+
+        self.assert_series_equal(ser, expected)
+
+    def test_setitem_nullable_boolean(self, data):
+        # GH 31446
+        ser = pd.Series([1] * len(data) + [0] * len(data), dtype="boolean")
+        ser[ser == 1] = 0
+        expected = pd.Series([0] * len(data) * 2, dtype="boolean")
+
+        self.assert_series_equal(ser, expected)
+
+    def test_setitem_nullable_string(self, data):
+        # GH 31446
+        ser = pd.Series(["a"] * len(data) + ["b"] * len(data), dtype="string")
+        ser[ser == "a"] = "c"
+        expected = pd.Series(["c"] * len(data) + ["b"] * len(data), dtype="string")
+
+        self.assert_series_equal(ser, expected)