feat: Add Collection.pipe method analogous to polars Lazy-/DataFrames

dataframely/collection/collection.py

@@ -8,11 +8,21 @@
 import textwrap
 import warnings
 from abc import ABC
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from dataclasses import asdict
 from json import JSONDecodeError
 from pathlib import Path
-from typing import IO, Annotated, Any, Literal, cast, overload
+from typing import (
+    IO,
+    Annotated,
+    Any,
+    Concatenate,
+    Literal,
+    ParamSpec,
+    TypeVar,
+    cast,
+    overload,
+)
 
 import polars as pl
 import polars.exceptions as plexc
@@ -53,6 +63,9 @@
 
 _FILTER_COLUMN_PREFIX = "__DATAFRAMELY_FILTER_COLUMN__"
 
+P = ParamSpec("P")
+T = TypeVar("T")
+
 
 class Collection(BaseCollection, ABC):
     """Base class for all collections of data frames with a predefined schema.
@@ -811,6 +824,33 @@ def collect_all(self) -> Self:
         dfs = pl.collect_all(lazy_dict.values())
         return self._init(dict(zip(lazy_dict, dfs)))
 
+    def pipe(
+        self,
+        function: Callable[Concatenate[Self, P], T],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> T:
+        """Apply a function to this collection.
+
+        This method allows chaining operations on a collection in a fluent style,
+        analogously to :meth:`polars.LazyFrame.pipe`.
+
+        Args:
+            function: The callable to apply. It receives this collection as its first
+                argument, followed by any additional ``args`` and ``kwargs``.
+            args: Additional positional arguments to pass to ``function``.
+            kwargs: Additional keyword arguments to pass to ``function``.
+
+        Returns:
+            The return value of ``function`` when called as described.
+
+        Example:
+            >>> def add_prefix(collection: MyCollection, prefix: str) -> MyCollection:
+            ...     ...
+            >>> result = my_collection.pipe(add_prefix, prefix="foo")
+        """
+        return function(self, *args, **kwargs)
+
     # --------------------------------- SERIALIZATION -------------------------------- #
 
     @classmethod

docs/api/collection/operations.rst

@@ -8,4 +8,5 @@ Operations
 
     Collection.collect_all
     Collection.join
+    Collection.pipe
     concat_collection_members

tests/collection/test_pipe.py

@@ -0,0 +1,55 @@
+# Copyright (c) QuantCo 2025-2026
+# SPDX-License-Identifier: BSD-3-Clause
+
+import dataframely as dy
+
+
+class SchemaOne(dy.Schema):
+    id = dy.Int64(primary_key=True)
+    name = dy.String(nullable=False)
+
+
+class SchemaTwo(dy.Schema):
+    id = dy.Int64(primary_key=True)
+    name = dy.String(nullable=False)
+
+
+class MyCollection(dy.Collection):
+    member_one: dy.LazyFrame[SchemaOne]
+    member_two: dy.LazyFrame[SchemaTwo]
+
+
+def test_pipe_passes_self() -> None:
+    # Arrange
+    collection = MyCollection.sample(overrides=[{"id": 1}, {"id": 2}])
+
+    # Act
+    result = collection.pipe(lambda c: c)
+
+    # Assert
+    assert result is collection
+
+
+def test_pipe_forwards_args_and_kwargs() -> None:
+    # Arrange
+    collection = MyCollection.sample(overrides=[{"id": 1}, {"id": 2}])
+
+    def combine(c: MyCollection, prefix: str, *, suffix: str) -> str:
+        return f"{prefix}{type(c).__name__}{suffix}"
+
+    # Act
+    result = collection.pipe(combine, "pre-", suffix="-post")
+
+    # Assert
+    assert result == "pre-MyCollection-post"
+
+
+def test_pipe_returns_arbitrary_type() -> None:
+    # Arrange
+    collection = MyCollection.sample(overrides=[{"id": 1}, {"id": 2}, {"id": 3}])
+
+    # Act
+    result = collection.pipe(lambda c: c.member_one.collect().height)
+
+    # Assert
+    assert result == 3