BUG/PERF: MaskedArray.__setitem__ validation (#45404)

Author: jbrockmendel

pandas/core/arrays/masked.py

@@ -55,6 +55,7 @@
 from pandas.core.dtypes.inference import is_array_like
 from pandas.core.dtypes.missing import (
     array_equivalent,
+    is_valid_na_for_dtype,
     isna,
     notna,
 )
@@ -233,17 +234,45 @@ def _coerce_to_array(
     ) -> tuple[np.ndarray, np.ndarray]:
         raise AbstractMethodError(cls)
 
-    def __setitem__(self, key, value) -> None:
-        _is_scalar = is_scalar(value)
-        if _is_scalar:
-            value = [value]
-        value, mask = self._coerce_to_array(value, dtype=self.dtype)
+    def _validate_setitem_value(self, value):
+        """
+        Check if we have a scalar that we can cast losslessly.
+
+        Raises
+        ------
+        TypeError
+        """
+        kind = self.dtype.kind
+        # TODO: get this all from np_can_hold_element?
+        if kind == "b":
+            if lib.is_bool(value):
+                return value
+
+        elif kind == "f":
+            if lib.is_integer(value) or lib.is_float(value):
+                return value
+
+        else:
+            if lib.is_integer(value) or (lib.is_float(value) and value.is_integer()):
+                return value
+            # TODO: unsigned checks
 
-        if _is_scalar:
-            value = value[0]
-            mask = mask[0]
+        raise TypeError(f"Invalid value '{value}' for dtype {self.dtype}")
 
+    def __setitem__(self, key, value) -> None:
         key = check_array_indexer(self, key)
+
+        if is_scalar(value):
+            if is_valid_na_for_dtype(value, self.dtype):
+                self._mask[key] = True
+            else:
+                value = self._validate_setitem_value(value)
+                self._data[key] = value
+                self._mask[key] = False
+            return
+
+        value, mask = self._coerce_to_array(value, dtype=self.dtype)
+
         self._data[key] = value
         self._mask[key] = mask
 

pandas/core/dtypes/missing.py

@@ -648,6 +648,9 @@ def is_valid_na_for_dtype(obj, dtype: DtypeObj) -> bool:
     elif dtype.kind in ["i", "u", "f", "c"]:
         # Numeric
         return obj is not NaT and not isinstance(obj, (np.datetime64, np.timedelta64))
+    elif dtype.kind == "b":
+        # We allow pd.NA, None, np.nan in BooleanArray (same as IntervalDtype)
+        return lib.is_float(obj) or obj is None or obj is libmissing.NA
 
     elif dtype == _dtype_str:
         # numpy string dtypes to avoid float np.nan

pandas/io/stata.py

@@ -593,7 +593,8 @@ def _cast_to_stata_types(data: DataFrame) -> DataFrame:
             missing_loc = data[col].isna()
             if missing_loc.any():
                 # Replace with always safe value
-                data.loc[missing_loc, col] = 0
+                fv = 0 if isinstance(data[col].dtype, _IntegerDtype) else False
+                data.loc[missing_loc, col] = fv
             # Replace with NumPy-compatible column
             data[col] = data[col].astype(data[col].dtype.numpy_dtype)
         dtype = data[col].dtype

pandas/tests/arrays/masked/test_indexing.py

@@ -0,0 +1,60 @@
+import re
+
+import numpy as np
+import pytest
+
+import pandas as pd
+
+
+class TestSetitemValidation:
+    def _check_setitem_invalid(self, arr, invalid):
+        msg = f"Invalid value '{str(invalid)}' for dtype {arr.dtype}"
+        msg = re.escape(msg)
+        with pytest.raises(TypeError, match=msg):
+            arr[0] = invalid
+
+        with pytest.raises(TypeError, match=msg):
+            arr[:] = invalid
+
+        with pytest.raises(TypeError, match=msg):
+            arr[[0]] = invalid
+
+        # FIXME: don't leave commented-out
+        # with pytest.raises(TypeError):
+        #    arr[[0]] = [invalid]
+
+        # with pytest.raises(TypeError):
+        #    arr[[0]] = np.array([invalid], dtype=object)
+
+        # Series non-coercion, behavior subject to change
+        ser = pd.Series(arr)
+        with pytest.raises(TypeError, match=msg):
+            ser[0] = invalid
+            # TODO: so, so many other variants of this...
+
+    _invalid_scalars = [
+        1 + 2j,
+        "True",
+        "1",
+        "1.0",
+        pd.NaT,
+        np.datetime64("NaT"),
+        np.timedelta64("NaT"),
+    ]
+
+    @pytest.mark.parametrize(
+        "invalid", _invalid_scalars + [1, 1.0, np.int64(1), np.float64(1)]
+    )
+    def test_setitem_validation_scalar_bool(self, invalid):
+        arr = pd.array([True, False, None], dtype="boolean")
+        self._check_setitem_invalid(arr, invalid)
+
+    @pytest.mark.parametrize("invalid", _invalid_scalars + [True, 1.5, np.float64(1.5)])
+    def test_setitem_validation_scalar_int(self, invalid, any_int_ea_dtype):
+        arr = pd.array([1, 2, None], dtype=any_int_ea_dtype)
+        self._check_setitem_invalid(arr, invalid)
+
+    @pytest.mark.parametrize("invalid", _invalid_scalars + [True])
+    def test_setitem_validation_scalar_float(self, invalid, float_ea_dtype):
+        arr = pd.array([1, 2, None], dtype=float_ea_dtype)
+        self._check_setitem_invalid(arr, invalid)

pandas/tests/frame/indexing/test_indexing.py

@@ -1240,12 +1240,10 @@ def test_setting_mismatched_na_into_nullable_fails(
 
         msg = "|".join(
             [
-                r"int\(\) argument must be a string, a bytes-like object or a "
-                "(real )?number, not 'NaTType'",
                 r"timedelta64\[ns\] cannot be converted to an? (Floating|Integer)Dtype",
                 r"datetime64\[ns\] cannot be converted to an? (Floating|Integer)Dtype",
-                "object cannot be converted to a FloatingDtype",
                 "'values' contains non-numeric NA",
+                r"Invalid value '.*' for dtype (U?Int|Float)\d{1,2}",
             ]
         )
         with pytest.raises(TypeError, match=msg):

pandas/tests/frame/indexing/test_where.py

@@ -907,16 +907,7 @@ def test_where_nullable_invalid_na(frame_or_series, any_numeric_ea_dtype):
 
     mask = np.array([True, True, False], ndmin=obj.ndim).T
 
-    msg = "|".join(
-        [
-            r"datetime64\[.{1,2}\] cannot be converted to an? (Integer|Floating)Dtype",
-            r"timedelta64\[.{1,2}\] cannot be converted to an? (Integer|Floating)Dtype",
-            r"int\(\) argument must be a string, a bytes-like object or a number, "
-            "not 'NaTType'",
-            "object cannot be converted to a FloatingDtype",
-            "'values' contains non-numeric NA",
-        ]
-    )
+    msg = r"Invalid value '.*' for dtype (U?Int|Float)\d{1,2}"
 
     for null in tm.NP_NAT_OBJECTS + [pd.NaT]:
         # NaT is an NA value that we should *not* cast to pd.NA dtype