BUG: membership checks on ExtensionArray containing NA values

Author: topper-123

doc/source/whatsnew/v1.2.0.rst

@@ -690,6 +690,7 @@ ExtensionArray
 - Fixed bug when applying a NumPy ufunc with multiple outputs to a :class:`pandas.arrays.IntegerArray` returning None (:issue:`36913`)
 - Fixed an inconsistency in :class:`PeriodArray`'s ``__init__`` signature to those of :class:`DatetimeArray` and :class:`TimedeltaArray` (:issue:`37289`)
 - Reductions for :class:`BooleanArray`, :class:`Categorical`, :class:`DatetimeArray`, :class:`FloatingArray`, :class:`IntegerArray`, :class:`PeriodArray`, :class:`TimedeltaArray`, and :class:`PandasArray` are now keyword-only methods (:issue:`37541`)
+- Bug, where a  `ValueError` was wrongly raised if a membership check was made on an `ExtensionArray` with :class:`NA` values, but without a custom ``__contains__`` method (:issue:`xxxxx`)
 
 Other
 ^^^^^

pandas/core/arrays/base.py

@@ -39,7 +39,7 @@
 )
 from pandas.core.dtypes.dtypes import ExtensionDtype
 from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
-from pandas.core.dtypes.missing import isna
+from pandas.core.dtypes.missing import isna, notna
 
 from pandas.core import ops
 from pandas.core.algorithms import factorize_array, unique
@@ -351,6 +351,18 @@ def __iter__(self):
         for i in range(len(self)):
             yield self[i]
 
+    def __contains__(self, item) -> bool:
+        """
+        Return for `item in self`.
+        """
+        # comparisons of ant item to pd.NA always return pd.NA, so e.g.
+        # "a" in [pd.NA] raises a TypeError. The implementation below works around that.
+        if isna(item):
+            return isna(self).any() if self._can_hold_na else False
+
+        arr = self[notna(self)] if self._can_hold_na else self
+        return item in iter(arr)
+
     def __eq__(self, other: Any) -> ArrayLike:
         """
         Return for `self == other` (element-wise equality).

pandas/tests/arrays/categorical/test_operators.py

@@ -395,3 +395,15 @@ def test_numeric_like_ops(self):
         msg = "Object with dtype category cannot perform the numpy op log"
         with pytest.raises(TypeError, match=msg):
             np.log(s)
+
+    def test_contains(self, ordered):
+        # GH-xxxxx
+        cat = Categorical(["a", "b"], ordered=ordered)
+        assert "a" in cat
+        assert "x" not in cat
+        assert pd.NA not in cat
+
+        cat = Categorical([np.nan, "a"], ordered=ordered)
+        assert "a" in cat
+        assert "x" not in cat
+        assert pd.NA in cat

pandas/tests/arrays/string_/test_string.py

@@ -345,3 +345,17 @@ def test_astype_from_float_dtype(dtype):
     result = s.astype("string")
     expected = pd.Series(["0.1"], dtype="string")
     tm.assert_series_equal(result, expected)
+
+
+def test_contains():
+    # GH-xxxxx
+    arr = pd.arrays.StringArray(np.array(["a", "b"], dtype=object))
+
+    assert "a" in arr
+    assert "x" not in arr
+    assert pd.NA not in arr
+
+    arr = pd.arrays.StringArray(np.array(["a", pd.NA]))
+    assert "a" in arr
+    assert "x" not in arr
+    assert pd.NA in arr