fix(python): Suggest str.contains for string containment in map_elements

Co-authored-by: Claude <noreply@anthropic.com>

Author: wtn

.github/scripts/test_bytecode_parser.py

@@ -21,7 +21,10 @@
 import pytest
 from polars._utils.udfs import BytecodeParser
 from tests.unit.operations.map.test_inefficient_map_warning import (
+    MY_COLLECTION,
     MY_DICT,
+    MY_STRING,
+    MY_SUBSTRING,
     NOOP_TEST_CASES,
     TEST_CASES,
 )
@@ -52,7 +55,10 @@ def test_bytecode_parser_expression_in_ipython(
         "from datetime import datetime; "
         "import numpy as np; "
         "import json; "
-        f"MY_DICT = {MY_DICT};"
+        f"MY_DICT = {MY_DICT}; "
+        f"MY_COLLECTION = {MY_COLLECTION}; "
+        f"MY_STRING = {repr(MY_STRING)}; "
+        f"MY_SUBSTRING = {repr(MY_SUBSTRING)}; "
         f'bytecode_parser = BytecodeParser({func}, map_target="expr");'
         f'print(bytecode_parser.to_expression("{col}"));'
     )

py-polars/src/polars/_utils/udfs.py

@@ -701,12 +701,44 @@ def _expr(self, value: StackEntry, col: str, param_name: str, depth: int) -> str
                     not_ = "" if op == "is" else "not_"
                     return f"{e1}.is_{not_}null()"
                 elif op in ("in", "not in"):
-                    not_ = "" if op == "in" else "~"
-                    return (
-                        f"{not_}({e1}.is_in({e2}))"
-                        if " " in e1
-                        else f"{not_}{e1}.is_in({e2})"
+                    e2_stripped = e2.lstrip()
+                    is_collection_literal = e2_stripped.startswith(
+                        ("(", "[", "{", "frozenset(")
                     )
+
+                    is_collection_variable = False
+                    if not is_collection_literal and not e2.startswith(
+                        ("pl.col(", "'")
+                    ):
+                        if not self._caller_variables:
+                            self._caller_variables = _get_all_caller_variables()
+                        var_value = self._caller_variables.get(e2)
+                        if isinstance(var_value, (list, tuple, set, frozenset, dict)):
+                            is_collection_variable = True
+
+                    if is_collection_literal or is_collection_variable:
+                        not_ = "" if op == "in" else "~"
+                        return (
+                            f"{not_}({e1}.is_in({e2}))"
+                            if " " in e1
+                            else f"{not_}{e1}.is_in({e2})"
+                        )
+                    else:
+                        e2_is_col = e2.startswith("pl.col(")
+                        e1_is_col = e1.startswith("pl.col(")
+
+                        if e2_is_col:
+                            needle = f"pl.lit({e1})" if not e1_is_col else e1
+                            haystack = e2
+                        else:
+                            needle = e1
+                            haystack = f"pl.lit({e2})"
+
+                        contains_expr = f"{haystack}.str.contains({needle})"
+
+                        if op == "not in":
+                            return f"~{contains_expr}"
+                        return contains_expr
                 elif op == "replace_strict":
                     if not self._caller_variables:
                         self._caller_variables = _get_all_caller_variables()

py-polars/tests/unit/operations/map/test_inefficient_map_warning.py

@@ -21,6 +21,9 @@
 MY_CONSTANT = 3
 MY_DICT = {0: "a", 1: "b", 2: "c", 3: "d", 4: "e"}
 MY_LIST = [1, 2, 3]
+MY_STRING = "qwerty"
+MY_SUBSTRING = "we"
+MY_COLLECTION = [2, 3, 4]
 
 # column_name, function, expected_suggestion
 TEST_CASES = [
@@ -67,12 +70,58 @@
     ),
     ("a", "lambda x: x in (2, 3, 4)", 'pl.col("a").is_in((2, 3, 4))', None),
     ("a", "lambda x: x not in (2, 3, 4)", '~pl.col("a").is_in((2, 3, 4))', None),
+    ("a", "lambda x: x in MY_COLLECTION", 'pl.col("a").is_in(MY_COLLECTION)', None),
+    ("a", "lambda x: x in MY_DICT", 'pl.col("a").is_in(MY_DICT)', None),
+    (
+        "a",
+        "lambda x: (x + 1) in (1, 2, 3)",
+        '((pl.col("a") + 1).is_in((1, 2, 3)))',
+        None,
+    ),
     (
         "a",
         "lambda x: x in (1, 2, 3, 4, 3) and x % 2 == 0 and x > 0",
         'pl.col("a").is_in((1, 2, 3, 4, 3)) & ((pl.col("a") % 2) == 0) & (pl.col("a") > 0)',
         None,
     ),
+    # ---------------------------------------------
+    # string containment with 'in' operator
+    # ---------------------------------------------
+    (
+        "b",
+        "lambda x: x in MY_STRING",
+        'pl.lit(MY_STRING).str.contains(pl.col("b"))',
+        None,
+    ),
+    (
+        "b",
+        "lambda x: MY_SUBSTRING in x",
+        'pl.col("b").str.contains(pl.lit(MY_SUBSTRING))',
+        None,
+    ),
+    ("b", 'lambda x: "A" in x', "pl.col(\"b\").str.contains(pl.lit('A'))", None),
+    (
+        "b",
+        "lambda x: x not in MY_STRING",
+        '~pl.lit(MY_STRING).str.contains(pl.col("b"))',
+        None,
+    ),
+    ("b", "lambda x: x in x", 'pl.col("b").str.contains(pl.col("b"))', None),
+    (
+        "b",
+        'lambda x: "test" in x',
+        "pl.col(\"b\").str.contains(pl.lit('test'))",
+        None,
+    ),
+    (
+        "b",
+        'lambda x: x not in "hello"',
+        "~pl.lit('hello').str.contains(pl.col(\"b\"))",
+        None,
+    ),
+    # ---------------------------------------------
+    # constants
+    # ---------------------------------------------
     ("a", "lambda x: MY_CONSTANT + x", 'MY_CONSTANT + pl.col("a")', None),
     (
         "a",
@@ -310,6 +359,9 @@
     "MY_CONSTANT": MY_CONSTANT,
     "MY_DICT": MY_DICT,
     "MY_LIST": MY_LIST,
+    "MY_STRING": MY_STRING,
+    "MY_SUBSTRING": MY_SUBSTRING,
+    "MY_COLLECTION": MY_COLLECTION,
     "cosh": cosh,
     "datetime": datetime,
     "dt": dt,