From 85257924956c69737b2b467d7195e0c29a1d473a Mon Sep 17 00:00:00 2001
From: Jeroen van Zundert <mail@jeroenvanzundert.nl>
Date: Fri, 13 Jan 2023 20:07:55 +0000
Subject: [PATCH] Synchronize input arg parsing between DataFrame and LazyFrame
 with_columns

---
 py-polars/polars/internals/dataframe/frame.py |  2 --
 py-polars/polars/internals/lazyframe/frame.py | 14 +++++++++-----
 py-polars/tests/unit/test_lazy.py             |  8 ++++++++
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/py-polars/polars/internals/dataframe/frame.py b/py-polars/polars/internals/dataframe/frame.py
index 7a233ac31783e..1edbc2cb592af 100644
--- a/py-polars/polars/internals/dataframe/frame.py
+++ b/py-polars/polars/internals/dataframe/frame.py
@@ -5571,8 +5571,6 @@ def with_columns(
         └─────┴──────┴───────┴──────┴───────┘
 
         """
-        if exprs is not None and not isinstance(exprs, Sequence):
-            exprs = [exprs]
         return (
             self.lazy().with_columns(exprs, **named_exprs).collect(no_optimization=True)
         )
diff --git a/py-polars/polars/internals/lazyframe/frame.py b/py-polars/polars/internals/lazyframe/frame.py
index 155f942dfd1e9..e5dc5923ac314 100644
--- a/py-polars/polars/internals/lazyframe/frame.py
+++ b/py-polars/polars/internals/lazyframe/frame.py
@@ -2474,11 +2474,15 @@ def with_columns(
         elif exprs is None and not named_exprs:
             raise ValueError("Expected at least one of 'exprs' or **named_exprs")
 
-        exprs = (
-            []
-            if exprs is None
-            else ([exprs] if isinstance(exprs, pli.Expr) else list(exprs))
-        )
+        if exprs is None:
+            exprs = []
+        elif isinstance(exprs, pli.Expr):
+            exprs = [exprs]
+        elif isinstance(exprs, pli.Series):
+            exprs = [pli.lit(exprs)]
+        else:
+            exprs = list(exprs)
+
         exprs.extend(
             pli.expr_to_lit_or_expr(expr).alias(name)
             for name, expr in named_exprs.items()
diff --git a/py-polars/tests/unit/test_lazy.py b/py-polars/tests/unit/test_lazy.py
index 384ad72dead7e..81a9e2ca20ddc 100644
--- a/py-polars/tests/unit/test_lazy.py
+++ b/py-polars/tests/unit/test_lazy.py
@@ -1000,6 +1000,14 @@ def test_with_column_renamed(fruits_cars: pl.DataFrame) -> None:
     assert res.columns[0] == "C"
 
 
+def test_with_columns_single_series() -> None:
+    df = pl.DataFrame({"a": [1, 2]})
+    result = df.lazy().with_columns(pl.Series("b", [3, 4])).collect()
+
+    expected = pl.DataFrame({"a": [1, 2], "b": [3, 4]})
+    assert_frame_equal(result, expected)
+
+
 def test_reverse() -> None:
     out = pl.DataFrame({"a": [1, 2], "b": [3, 4]}).lazy().reverse()
     expected = pl.DataFrame({"a": [2, 1], "b": [4, 3]})