Add Series to with_columns() argument annotations (#3696)

Co-authored-by: Adam Gregory <adam.gregory@diasemi.com>
pola-rs · Jun 17, 2022 · af3327e · af3327e
1 parent 25e9e35
commit af3327e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/py-polars/polars/internals/frame.py b/py-polars/polars/internals/frame.py
@@ -4696,7 +4696,10 @@ def select(
             .collect(no_optimization=True, string_cache=False)
         )
 
-    def with_columns(self: DF, exprs: Union["pli.Expr", List["pli.Expr"]]) -> DF:
+    def with_columns(
+        self: DF,
+        exprs: Union["pli.Expr", "pli.Series", List[Union["pli.Expr", "pli.Series"]]],
+    ) -> DF:
         """
         Add or overwrite multiple columns in a DataFrame.
 

diff --git a/py-polars/polars/internals/lazy_frame.py b/py-polars/polars/internals/lazy_frame.py
@@ -1430,7 +1430,10 @@ def join(
             )
         )
 
-    def with_columns(self: LDF, exprs: Union[List["pli.Expr"], "pli.Expr"]) -> LDF:
+    def with_columns(
+        self: LDF,
+        exprs: Union["pli.Expr", "pli.Series", List[Union["pli.Expr", "pli.Series"]]],
+    ) -> LDF:
         """
         Add or overwrite multiple columns in a DataFrame.