diff --git a/dataframely/collection/collection.py b/dataframely/collection/collection.py index 43aab21c..059d52b6 100644 --- a/dataframely/collection/collection.py +++ b/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 diff --git a/docs/api/collection/operations.rst b/docs/api/collection/operations.rst index 46b1000b..8026c7ec 100644 --- a/docs/api/collection/operations.rst +++ b/docs/api/collection/operations.rst @@ -8,4 +8,5 @@ Operations Collection.collect_all Collection.join + Collection.pipe concat_collection_members diff --git a/tests/collection/test_pipe.py b/tests/collection/test_pipe.py new file mode 100644 index 00000000..8d66d486 --- /dev/null +++ b/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