ENH: Enable indexing with nullable Boolean

pandas/tests/extension/base/setitem.py

@@ -98,8 +98,9 @@ def test_setitem_iloc_scalar_multiple_homogoneous(self, data):
         [
             np.array([True, True, True, False, False]),
             pd.array([True, True, True, False, False], dtype="boolean"),
+            pd.array([True, True, True, pd.NA, pd.NA], dtype="boolean"),
         ],
-        ids=["numpy-array", "boolean-array"],
+        ids=["numpy-array", "boolean-array", "boolean-array-na"],
     )
     def test_setitem_mask(self, data, mask, box_in_series):
         arr = data[:5].copy()
@@ -124,20 +125,17 @@ def test_setitem_mask_raises(self, data, box_in_series):
         with pytest.raises(IndexError, match="wrong length"):
             data[mask] = data[0]
 
-    def test_setitem_mask_boolean_array_raises(self, data, box_in_series):
-        # missing values in mask
+    def test_setitem_mask_boolean_array_with_na(self, data, box_in_series):
         mask = pd.array(np.zeros(data.shape, dtype="bool"), dtype="boolean")
-        mask[:2] = pd.NA
+        mask[:3] = True
+        mask[3:5] = pd.NA
 
         if box_in_series:
             data = pd.Series(data)
 
-        msg = (
-            "Cannot mask with a boolean indexer containing NA values|"
-            "cannot mask with array containing NA / NaN values"
-        )
-        with pytest.raises(ValueError, match=msg):
-            data[mask] = data[0]
+        data[mask] = data[0]
+
+        assert (data[:3] == data[0]).all()
 
     @pytest.mark.parametrize(
         "idx",
@@ -219,36 +217,6 @@ def test_setitem_mask_broadcast(self, data, setter):
         assert ser[0] == data[10]
         assert ser[1] == data[10]
 
-    def test_setitem_boolean_na_mask(self, data, box_in_series):
-        # https://github.com/pandas-dev/pandas/issues/31503
-        if box_in_series:
-            data = pd.Series(data)
-
-        mask = pd.array(np.zeros(len(data), dtype=bool), dtype="boolean")
-        mask[:2] = pd.NA
-        mask[2:4] = True
-        original = data.copy()
-
-        data[mask] = data[mask.fillna(False)]
-
-        if box_in_series:
-            tm.assert_series_equal(data, original)
-        else:
-            tm.assert_extension_array_equal(data, original)
-
-    def test_setitem_boolean_na_mask_frame(self, data):
-        # https://github.com/pandas-dev/pandas/issues/31503
-        df = pd.DataFrame({"a": range(len(data)), "b": data})
-        original = df.copy()
-
-        mask = pd.array(np.zeros(len(data), dtype=bool), dtype="boolean")
-        mask[:2] = pd.NA
-        mask[2:4] = True
-
-        df.loc[mask, "b"] = data[mask.fillna(False)]
-
-        tm.assert_frame_equal(df, original)
-
     def test_setitem_expand_columns(self, data):
         df = pd.DataFrame({"A": data})
         result = df.copy()

pandas/tests/extension/test_numpy.py

@@ -415,10 +415,6 @@ def test_setitem_mask(self, data, mask, box_in_series):
     def test_setitem_mask_raises(self, data, box_in_series):
         super().test_setitem_mask_raises(data, box_in_series)
 
-    @skip_nested
-    def test_setitem_mask_boolean_array_raises(self, data, box_in_series):
-        super().test_setitem_mask_boolean_array_raises(data, box_in_series)
-
     @skip_nested
     @pytest.mark.parametrize(
         "idx",