feat(python): Allow iterable of frames as input to align_frames (#2…

…1209)

py-polars/polars/_utils/deprecation.py

@@ -211,7 +211,7 @@ def deprecate_parameter_as_multi_positional(
 
     Use as follows::
 
-        @deprecate_parameter_as_positional("columns", version="0.20.4")
+        @deprecate_parameter_as_multi_positional("columns", version="0.20.4")
         def myfunc(*columns): ...
     """  # noqa: W505
 

py-polars/polars/functions/eager.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import contextlib
-from collections.abc import Sequence
+from collections.abc import Generator, Iterator, Sequence
 from functools import reduce
 from itertools import chain
 from typing import TYPE_CHECKING, get_args
@@ -337,7 +337,7 @@ def join_func(
 
 
 def align_frames(
-    *frames: FrameType,
+    *frames: FrameType | Iterable[FrameType],
     on: str | Expr | Sequence[str] | Sequence[Expr] | Sequence[str | Expr],
     how: JoinStrategy = "full",
     select: str | Expr | Sequence[str | Expr] | None = None,
@@ -479,6 +479,11 @@ def align_frames(
     if not frames:
         return []
 
+    if len(frames) == 1 and not isinstance(frames[0], (pl.DataFrame, pl.LazyFrame)):
+        frames = frames[0]  # type: ignore[assignment]
+    if isinstance(frames, (Generator, Iterator)):
+        frames = tuple(frames)
+
     if len({type(f) for f in frames}) != 1:
         msg = (
             "input frames must be of a consistent type (all LazyFrame or all DataFrame)"
@@ -489,9 +494,9 @@ def align_frames(
     on = [on] if (isinstance(on, str) or not isinstance(on, Sequence)) else on
     align_on = [(c.meta.output_name() if isinstance(c, pl.Expr) else c) for c in on]
 
-    # create aligned master frame (this is the most expensive part; afterwards
-    # we just subselect out the columns representing the component frames)
-    idx_frames = [(idx, frame.lazy()) for idx, frame in enumerate(frames)]
+    # create aligned master frame (this is the most expensive part; after
+    # we just select out the columns representing the component frames)
+    idx_frames = [(idx, frame.lazy()) for idx, frame in enumerate(frames)]  # type: ignore[union-attr]
     alignment_frame = _alignment_join(
         *idx_frames, align_on=align_on, how=how, descending=descending
     )

py-polars/tests/unit/functions/test_functions.py

@@ -310,16 +310,24 @@ def test_align_frames() -> None:
 
 
 def test_align_frames_misc() -> None:
-    # descending result
     df1 = pl.DataFrame([[3, 5, 6], [5, 8, 9]], orient="row")
     df2 = pl.DataFrame([[2, 5, 6], [3, 8, 9], [4, 2, 0]], orient="row")
 
-    pf1, pf2 = pl.align_frames(df1, df2, on="column_0", descending=True)
+    # descending result
+    pf1, pf2 = pl.align_frames(
+        [df1, df2],  # list input
+        on="column_0",
+        descending=True,
+    )
     assert pf1.rows() == [(5, 8, 9), (4, None, None), (3, 5, 6), (2, None, None)]
     assert pf2.rows() == [(5, None, None), (4, 2, 0), (3, 8, 9), (2, 5, 6)]
 
     # handle identical frames
-    pf1, pf2, pf3 = pl.align_frames(df1, df2, df2, on="column_0", descending=True)
+    pf1, pf2, pf3 = pl.align_frames(
+        (df for df in (df1, df2, df2)),  # generator input
+        on="column_0",
+        descending=True,
+    )
     assert pf1.rows() == [(5, 8, 9), (4, None, None), (3, 5, 6), (2, None, None)]
     for pf in (pf2, pf3):
         assert pf.rows() == [(5, None, None), (4, 2, 0), (3, 8, 9), (2, 5, 6)]