fix(python): Reject Expr objects in replace sequence arguments

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

Author: wtn

py-polars/src/polars/expr/expr.py

@@ -11165,16 +11165,38 @@ def replace(
                 raise TypeError(msg)
             new = list(old.values())
             old = list(old.keys())
+
+        old_is_seq = isinstance(old, Sequence) and not isinstance(old, (str, pl.Series))
+        new_is_seq = isinstance(new, Sequence) and not isinstance(new, (str, pl.Series))
+        has_expr_in_old = old_is_seq and any(isinstance(v, Expr) for v in old)  # type: ignore[union-attr]
+        has_expr_in_new = new_is_seq and any(isinstance(v, Expr) for v in new)  # type: ignore[union-attr]
+
+        if has_expr_in_old or has_expr_in_new:
+            old_list = list(old) if old_is_seq else [old]  # type: ignore[arg-type, misc]
+            new_list = list(new) if new_is_seq else [new]  # type: ignore[arg-type]
+
+            if len(new_list) == 1 and len(old_list) > 1:
+                new_list = new_list * len(old_list)
+
+            if len(old_list) != len(new_list):
+                msg = f"lengths of `old` ({len(old_list)}) and `new` ({len(new_list)}) must match"
+                raise ValueError(msg)
+
+            # when(False).then(self) preserves the column name
+            result: Expr = F.when(F.lit(False)).then(self)
+            for old_val, new_val in zip(old_list, new_list, strict=True):
+                result = result.when(self == old_val).then(new_val)  # type: ignore[attr-defined]
+            result = result.otherwise(self)  # type: ignore[attr-defined]
         else:
-            if isinstance(old, Sequence) and not isinstance(old, (str, pl.Series)):
+            if old_is_seq:
                 old = pl.Series(old)
-            if isinstance(new, Sequence) and not isinstance(new, (str, pl.Series)):
+            if new_is_seq:
                 new = pl.Series(new)
 
-        old_pyexpr = parse_into_expression(old, str_as_lit=True)  # type: ignore[arg-type]
-        new_pyexpr = parse_into_expression(new, str_as_lit=True)
+            old_pyexpr = parse_into_expression(old, str_as_lit=True)  # type: ignore[arg-type]
+            new_pyexpr = parse_into_expression(new, str_as_lit=True)  # type: ignore[arg-type]
 
-        result = wrap_expr(self._pyexpr.replace(old_pyexpr, new_pyexpr))
+            result = wrap_expr(self._pyexpr.replace(old_pyexpr, new_pyexpr))
 
         if return_dtype is not None:
             result = result.cast(return_dtype)
@@ -11362,6 +11384,37 @@ def replace_strict(
             new = list(old.values())
             old = list(old.keys())
 
+        old_is_seq = isinstance(old, Sequence) and not isinstance(old, (str, pl.Series))
+        new_is_seq = isinstance(new, Sequence) and not isinstance(new, (str, pl.Series))
+        has_expr_in_old = old_is_seq and any(isinstance(v, Expr) for v in old)  # type: ignore[union-attr]
+        has_expr_in_new = new_is_seq and any(isinstance(v, Expr) for v in new)  # type: ignore[union-attr]
+
+        if has_expr_in_old or has_expr_in_new:
+            old_list = list(old) if old_is_seq else [old]  # type: ignore[arg-type, misc]
+            new_list = list(new) if new_is_seq else [new]  # type: ignore[arg-type]
+
+            if len(new_list) == 1 and len(old_list) > 1:
+                new_list = new_list * len(old_list)
+
+            if len(old_list) != len(new_list):
+                msg = f"lengths of `old` ({len(old_list)}) and `new` ({len(new_list)}) must match"
+                raise ValueError(msg)
+
+            # when(False).then(self) preserves the column name
+            result: Expr = F.when(F.lit(False)).then(self)
+            for old_val, new_val in zip(old_list, new_list, strict=True):
+                result = result.when(self == old_val).then(new_val)  # type: ignore[attr-defined]
+
+            if default is no_default:
+                result = result.otherwise(None)  # type: ignore[attr-defined]
+            else:
+                result = result.otherwise(default)  # type: ignore[attr-defined]
+
+            if return_dtype is not None:
+                result = result.cast(return_dtype)
+
+            return result
+
         old_pyexpr = parse_into_expression(old, str_as_lit=True)  # type: ignore[arg-type]
         new_pyexpr = parse_into_expression(new, str_as_lit=True)  # type: ignore[arg-type]
 

py-polars/tests/unit/operations/test_replace.py

@@ -289,3 +289,54 @@ def test_replace_single_argument_not_mapping() -> None:
         match="`new` argument is required if `old` argument is not a Mapping type",
     ):
         df.select(pl.col("a").replace("b"))
+
+
+def test_replace_expr_in_sequence() -> None:
+    s = pl.Series([1, 2, 3])
+
+    result = s.replace([pl.lit(1)], [10])
+    expected = pl.Series([10, 2, 3])
+    assert_series_equal(result, expected)
+
+    result = s.replace([1], [pl.lit(10)])
+    expected = pl.Series([10, 2, 3])
+    assert_series_equal(result, expected)
+
+    result = s.replace([pl.lit(1), 2], [10, pl.lit(20)])
+    expected = pl.Series([10, 20, 3])
+    assert_series_equal(result, expected)
+
+    result = s.replace(pl.lit(1), pl.lit(10))
+    expected = pl.Series([10, 2, 3])
+    assert_series_equal(result, expected)
+
+
+def test_replace_expr_in_sequence_with_column_refs() -> None:
+    df = pl.DataFrame({"a": [1, 2, 3], "b": [10, 20, 30]})
+
+    result = df.select(
+        pl.col("a").replace([pl.col("a").min(), pl.col("a").max()], [100, 300])
+    )
+    expected = pl.DataFrame({"a": [100, 2, 300]})
+    assert_frame_equal(result, expected)
+
+    result = df.select(
+        pl.col("a").replace([1, 2], [pl.col("b").first(), pl.col("b").sum()])
+    )
+    expected = pl.DataFrame({"a": [10, 60, 3]})
+    assert_frame_equal(result, expected)
+
+
+def test_replace_expr_in_sequence_many_to_one() -> None:
+    s = pl.Series([1, 2, 3])
+    result = s.replace([pl.lit(1), 2], [10])
+    expected = pl.Series([10, 10, 3])
+    assert_series_equal(result, expected)
+
+
+def test_replace_expr_in_sequence_length_mismatch() -> None:
+    s = pl.Series([1, 2, 3])
+    with pytest.raises(
+        ValueError, match=r"lengths of `old` \(2\) and `new` \(3\) must match"
+    ):
+        s.replace([pl.lit(1), 2], [10, 20, 30])

py-polars/tests/unit/operations/test_replace_strict.py

@@ -419,3 +419,52 @@ def test_replace_strict_nested_mapping_22554() -> None:
         ),
         pl.Series([[42], [13], [37]]),
     )
+
+
+def test_replace_strict_expr_in_sequence() -> None:
+    s = pl.Series([1, 2, 3])
+
+    result = s.replace_strict([pl.lit(1)], [10], default=None)
+    expected = pl.Series([10, None, None])
+    assert_series_equal(result, expected, check_dtypes=False)
+
+    result = s.replace_strict([1], [pl.lit(10)], default=None)
+    expected = pl.Series([10, None, None])
+    assert_series_equal(result, expected, check_dtypes=False)
+
+    result = s.replace_strict([pl.lit(1), 2], [10, pl.lit(20)], default=-1)
+    expected = pl.Series([10, 20, -1])
+    assert_series_equal(result, expected, check_dtypes=False)
+
+    result = s.replace_strict(pl.lit(1), pl.lit(10), default=pl.lit(None))
+    expected = pl.Series([10, None, None], dtype=pl.Int32)
+    assert_series_equal(result, expected)
+
+
+def test_replace_strict_expr_in_sequence_many_to_one() -> None:
+    s = pl.Series([1, 2, 3])
+    result = s.replace_strict([pl.lit(1), 2], [10], default=-1)
+    expected = pl.Series([10, 10, -1])
+    assert_series_equal(result, expected, check_dtypes=False)
+
+
+def test_replace_strict_expr_in_sequence_length_mismatch() -> None:
+    s = pl.Series([1, 2, 3])
+    with pytest.raises(
+        ValueError, match=r"lengths of `old` \(2\) and `new` \(3\) must match"
+    ):
+        s.replace_strict([pl.lit(1), 2], [10, 20, 30], default=None)
+
+
+def test_replace_strict_expr_in_sequence_no_default() -> None:
+    s = pl.Series([1, 2, 3])
+    result = s.replace_strict([pl.lit(1)], [10])
+    expected = pl.Series([10, None, None])
+    assert_series_equal(result, expected, check_dtypes=False)
+
+
+def test_replace_strict_expr_in_sequence_return_dtype() -> None:
+    s = pl.Series([1, 2, 3])
+    result = s.replace_strict([pl.lit(1)], [10], default=None, return_dtype=pl.Float64)
+    expected = pl.Series([10.0, None, None], dtype=pl.Float64)
+    assert_series_equal(result, expected)