feat: implement __contains__ for DataFrame and LazyFrame (#4035)

py-polars/polars/internals/frame.py

@@ -1598,6 +1598,9 @@ def __getattr__(self, item: Any) -> PySeries:
         except Exception:
             raise AttributeError(item)
 
+    def __contains__(self, key: str) -> bool:
+        return key in self.columns
+
     def __iter__(self) -> Iterator[Any]:
         return self.get_columns().__iter__()
 

py-polars/polars/internals/lazy_frame.py

@@ -360,6 +360,9 @@ def __getitem__(self: LDF, item: int | range | slice) -> LazyFrame:
             )
         return LazyPolarsSlice(self).apply(item)
 
+    def __contains__(self: LDF, key: str) -> bool:
+        return key in self.columns
+
     def pipe(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
         """
         Apply a function on Self.

py-polars/tests/test_df.py

@@ -509,6 +509,13 @@ def test_from_arrow() -> None:
     assert pl.from_arrow(tbl).shape == (2, 5)
 
 
+def test_dataframe_membership_operator() -> None:
+    # cf. issue #4032
+    df = pl.DataFrame({"name": ["Jane", "John"], "age": [20, 30]})
+    assert "name" in df
+    assert "phone" not in df
+
+
 def test_sort() -> None:
     df = pl.DataFrame({"a": [2, 1, 3], "b": [1, 2, 3]})
     with pytest.deprecated_call():

py-polars/tests/test_lazy.py

@@ -26,6 +26,12 @@ def test_lazy() -> None:
     df.groupby("a").agg(pl.list("b"))
 
 
+def test_lazyframe_membership_operator() -> None:
+    ldf = pl.DataFrame({"name": ["Jane", "John"], "age": [20, 30]}).lazy()
+    assert "name" in ldf
+    assert "phone" not in ldf
+
+
 def test_apply() -> None:
     df = pl.DataFrame({"a": [1, 2, 3], "b": [1.0, 2.0, 3.0]})
     new = df.lazy().with_column(col("a").map(lambda s: s * 2).alias("foo")).collect()