API: membership checks on ExtensionArray containing NA values

doc/source/whatsnew/v1.2.0.rst

@@ -752,6 +752,7 @@ ExtensionArray
 - Fixed bug when applying a NumPy ufunc with multiple outputs to an :class:`.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`)
+- Fixed a bug where a  ``ValueError`` was wrongly raised if a membership check was made on an ``ExtensionArray`` containing nan-like values (:issue:`37867`)
 
 Other
 ^^^^^

pandas/core/arrays/base.py

@@ -37,6 +37,7 @@
     is_array_like,
     is_dtype_equal,
     is_list_like,
+    is_scalar,
     pandas_dtype,
 )
 from pandas.core.dtypes.dtypes import ExtensionDtype
@@ -354,6 +355,19 @@ def __iter__(self):
         for i in range(len(self)):
             yield self[i]
 
+    def __contains__(self, item) -> bool:
+        """
+        Return for `item in self`.
+        """
+        # comparisons of any item to pd.NA always return pd.NA, so e.g. "a" in [pd.NA]
+        # would raise a TypeError. The implementation below works around that.
+        if item is self.dtype.na_value:
+            return self.isna().any() if self._can_hold_na else False
+        elif is_scalar(item) and isna(item):
+            return False
+        else:
+            return (item == self).any()
+
     def __eq__(self, other: Any) -> ArrayLike:
         """
         Return for `self == other` (element-wise equality).

pandas/tests/arrays/categorical/test_operators.py

@@ -395,3 +395,23 @@ 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-37867
+        cat = Categorical(["a", "b"], ordered=ordered)
+        assert "a" in cat
+        assert "x" not in cat
+        assert np.nan 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 np.nan in cat
+        assert pd.NA in cat
+
+        cat = cat[::-1]
+        assert "a" in cat
+        assert "x" not in cat
+        assert np.nan in cat
+        assert pd.NA in cat

pandas/tests/arrays/string_/test_string.py

@@ -524,3 +524,19 @@ def test_to_numpy_na_value(dtype, nulls_fixture):
     result = arr.to_numpy(na_value=na_value)
     expected = np.array(["a", na_value, "b"], dtype=object)
     tm.assert_numpy_array_equal(result, expected)
+
+
+def test_contains():
+    # GH-37867
+    arr = pd.array(["a", "b"], dtype="string")
+
+    assert "a" in arr
+    assert "x" not in arr
+    assert np.nan not in arr
+    assert pd.NA not in arr
+
+    arr = pd.array(["a", pd.NA], dtype="string")
+    assert "a" in arr
+    assert "x" not in arr
+    assert np.nan not in arr
+    assert pd.NA in arr

pandas/tests/extension/base/interface.py

@@ -29,6 +29,31 @@ def test_can_hold_na_valid(self, data):
         # GH-20761
         assert data._can_hold_na is True
 
+    def test_contains(self, data, data_missing):
+        # GH-37867
+        # Tests for membership checks. Membership checks for nan-likes is tricky and
+        # the settled on rule is: `nan_like in arr` is True if nan_like is
+        # arr.dtype.na_value and arr.isna().any() is True. Else the check returns False.
+
+        for this_data in [data, data_missing]:
+            scalar = this_data[~this_data.isna()][0]
+
+            assert scalar in this_data
+            assert "124jhujbhjhb5" not in data
+
+            na_value = this_data.dtype.na_value
+
+            if this_data.isna().any():
+                assert na_value in this_data
+            else:
+                assert na_value not in this_data
+
+            # this_data can never contain other nan-likes than na_value
+            for na_value_type in {None, np.nan, pd.NA, pd.NaT}:
+                if na_value_type is na_value:
+                    continue
+                assert na_value_type not in this_data
+
     def test_memory_usage(self, data):
         s = pd.Series(data)
         result = s.memory_usage(index=False)

pandas/tests/extension/json/test_json.py

@@ -143,6 +143,11 @@ def test_custom_asserts(self):
         with pytest.raises(AssertionError, match=msg):
             self.assert_frame_equal(a.to_frame(), b.to_frame())
 
+    @pytest.mark.xfail(reason="comparison method not implemented on JSONArray")
+    def test_contains(self, data):
+        # GH-37867
+        super().test_contains(data)
+
 
 class TestConstructors(BaseJSON, base.BaseConstructorsTests):
     @pytest.mark.skip(reason="not implemented constructor from dtype")