feat: add Series.__contains__

narwhals/__init__.py

@@ -31,22 +31,22 @@
 from narwhals.dtypes import UInt64
 from narwhals.dtypes import Unknown
 from narwhals.expr import Expr
-from narwhals.expr import all_ as all  # noqa: A004
+from narwhals.expr import all_ as all
 from narwhals.expr import all_horizontal
 from narwhals.expr import any_horizontal
 from narwhals.expr import col
 from narwhals.expr import concat_str
-from narwhals.expr import len_ as len  # noqa: A004
+from narwhals.expr import len_ as len
 from narwhals.expr import lit
-from narwhals.expr import max  # noqa: A004
+from narwhals.expr import max
 from narwhals.expr import max_horizontal
 from narwhals.expr import mean
 from narwhals.expr import mean_horizontal
 from narwhals.expr import median
-from narwhals.expr import min  # noqa: A004
+from narwhals.expr import min
 from narwhals.expr import min_horizontal
 from narwhals.expr import nth
-from narwhals.expr import sum  # noqa: A004
+from narwhals.expr import sum
 from narwhals.expr import sum_horizontal
 from narwhals.expr import when
 from narwhals.functions import concat

narwhals/_arrow/series.py

@@ -1021,6 +1021,31 @@ def __iter__(self: Self) -> Iterator[Any]:
             for x in self._native_series.__iter__()
         )
 
+    def __contains__(self: Self, other: Any) -> bool:
+        from pyarrow import ArrowInvalid  # ignore-banned-imports
+        from pyarrow import ArrowNotImplementedError  # ignore-banned-imports
+        from pyarrow import ArrowTypeError  # ignore-banned-imports
+
+        try:
+            import pyarrow as pa  # ignore-banned-imports
+            import pyarrow.compute as pc  # ignore-banned-imports
+
+            native_series = self._native_series
+            other_ = (
+                pa.scalar(other)
+                if other is not None
+                else pa.scalar(None, type=native_series.type)
+            )
+            return maybe_extract_py_scalar(  # type: ignore[no-any-return]
+                pc.is_in(other_, native_series),
+                return_py_scalar=True,
+            )
+        except (ArrowInvalid, ArrowNotImplementedError, ArrowTypeError) as exc:
+            from narwhals.exceptions import InvalidOperationError
+
+            msg = f"Unable to compare other of type {type(other)} with series of type {self.dtype}."
+            raise InvalidOperationError(msg) from exc
+
     @property
     def shape(self: Self) -> tuple[int]:
         return (len(self._native_series),)

narwhals/_pandas_like/series.py

@@ -841,6 +841,13 @@ def rolling_mean(
     def __iter__(self: Self) -> Iterator[Any]:
         yield from self._native_series.__iter__()
 
+    def __contains__(self: Self, other: Any) -> bool:
+        return (  # type: ignore[no-any-return]
+            self._native_series.isna().any()
+            if other is None
+            else (self._native_series == other).any()
+        )
+
     def is_finite(self: Self) -> Self:
         s = self._native_series
         return self._from_native_series((s > float("-inf")) & (s < float("inf")))

narwhals/_polars/series.py

@@ -343,6 +343,17 @@ def cum_count(self: Self, *, reverse: bool) -> Self:
 
         return self._from_native_series(result)
 
+    def __contains__(self: Self, other: Any) -> bool:
+        from polars.exceptions import InvalidOperationError as PlInvalidOperationError
+
+        try:
+            return self._native_series.__contains__(other)
+        except PlInvalidOperationError as exc:
+            from narwhals.exceptions import InvalidOperationError
+
+            msg = f"Unable to compare other of type {type(other)} with series of type {self.dtype}."
+            raise InvalidOperationError(msg) from exc
+
     @property
     def dt(self: Self) -> PolarsSeriesDateTimeNamespace:
         return PolarsSeriesDateTimeNamespace(self)

narwhals/series.py

@@ -3324,6 +3324,9 @@ def rolling_mean(
     def __iter__(self: Self) -> Iterator[Any]:
         yield from self._compliant_series.__iter__()
 
+    def __contains__(self: Self, other: Any) -> bool:
+        return self._compliant_series.__contains__(other)  # type: ignore[no-any-return]
+
     @property
     def str(self: Self) -> SeriesStringNamespace[Self]:
         return SeriesStringNamespace(self)

narwhals/stable/v1/selectors.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from narwhals.selectors import all  # noqa: A004
+from narwhals.selectors import all
 from narwhals.selectors import boolean
 from narwhals.selectors import by_dtype
 from narwhals.selectors import categorical

pyproject.toml

@@ -53,6 +53,7 @@ lint.select = [
 ]
 lint.ignore = [
   "A001",
+  "A004",
   "ARG002",
   "ANN401",
   "C901",

tests/selectors_test.py

@@ -5,7 +5,7 @@
 import pytest
 
 import narwhals.stable.v1 as nw
-from narwhals.stable.v1.selectors import all  # noqa: A004
+from narwhals.stable.v1.selectors import all
 from narwhals.stable.v1.selectors import boolean
 from narwhals.stable.v1.selectors import by_dtype
 from narwhals.stable.v1.selectors import categorical

tests/series_only/__contains___test.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from typing import Any
+
+import pytest
+
+import narwhals.stable.v1 as nw
+from narwhals.exceptions import InvalidOperationError
+
+if TYPE_CHECKING:
+    from tests.utils import ConstructorEager
+
+data = [100, 200, None]
+
+
+@pytest.mark.parametrize(
+    ("other", "expected"), [(100, True), (None, True), (1, False), (100.314, False)]
+)
+def test_contains(
+    constructor_eager: ConstructorEager,
+    other: int | None,
+    expected: bool,  # noqa: FBT001
+) -> None:
+    s = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"]
+
+    assert (other in s) == expected
+
+
+@pytest.mark.parametrize("other", ["foo", [1, 2, 3]])
+def test_contains_invalid_type(
+    request: pytest.FixtureRequest,
+    constructor_eager: ConstructorEager,
+    other: Any,
+) -> None:
+    if "polars" not in str(constructor_eager) and "pyarrow_table" not in str(
+        constructor_eager
+    ):
+        request.applymarker(pytest.mark.xfail)
+
+    s = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"]
+
+    with pytest.raises(InvalidOperationError):
+        _ = other in s

utils/check_api_reference.py

@@ -32,6 +32,7 @@
     "value_counts",
     "zip_with",
     "__iter__",
+    "__contains__",
 }
 BASE_DTYPES = {
     "NumericType",