Repository: narwhals-dev/narwhals
Pull request #115 (2024-05-08 16:42:56)

add pipe


Patch:
diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md
index 4cb39bfa..3f807456 100644
--- a/docs/api-reference/dataframe.md
+++ b/docs/api-reference/dataframe.md
@@ -11,6 +11,7 @@
         - head
         - join
         - lazy
+        - pipe
         - rename
         - schema
         - select
diff --git a/docs/api-reference/lazyframe.md b/docs/api-reference/lazyframe.md
index 811c1785..7f858b34 100644
--- a/docs/api-reference/lazyframe.md
+++ b/docs/api-reference/lazyframe.md
@@ -12,6 +12,7 @@
         - head
         - join
         - lazy
+        - pipe
         - rename
         - schema
         - select
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 3de3cc6e..74736447 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import Callable
 from typing import Iterable
 from typing import Literal
 from typing import Sequence
@@ -77,6 +78,9 @@ class BaseFrame:
             for k, v in self._dataframe.schema.items()
         }
 
+    def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Self:
+        return function(self, *args, **kwargs)
+
     @property
     def columns(self) -> list[str]:
         return self._dataframe.columns  # type: ignore[no-any-return]
@@ -433,6 +437,46 @@ class DataFrame(BaseFrame):
         return self._dataframe.to_dict(as_series=as_series)  # type: ignore[no-any-return]
 
     # inherited
+    def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Self:
+        """
+        Pipe function call.
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> data = {'a': [1,2,3], 'ba': [4,5,6]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            Let's define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...     df = nw.from_native(df_any)
+            ...     df = df.pipe(lambda _df: _df.select([x for x in _df.columns if len(x) == 1]))
+            ...     return nw.to_native(df)
+
+            We can then pass either pandas or Polars:
+
+            >>> func(df_pd)
+               a
+            0  1
+            1  2
+            2  3
+            >>> func(df_pl)
+            shape: (3, 1)
+            ┌─────┐
+            │ a   │
+            │ --- │
+            │ i64 │
+            ╞═════╡
+            │ 1   │
+            │ 2   │
+            │ 3   │
+            └─────┘
+        """
+        return super().pipe(function, *args, **kwargs)
+
     @property
     def schema(self) -> dict[str, DType]:
         r"""
@@ -1342,6 +1386,46 @@ class LazyFrame(BaseFrame):
         )
 
     # inherited
+    def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Self:
+        """
+        Pipe function call.
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> data = {'a': [1,2,3], 'ba': [4,5,6]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.LazyFrame(data)
+
+            Let's define a dataframe-agnostic function:
+
+            >>> def func(df_any):
+            ...     df = nw.from_native(df_any)
+            ...     df = df.pipe(lambda _df: _df.select([x for x in _df.columns if len(x) == 1]))
+            ...     return nw.to_native(df)
+
+            We can then pass either pandas or Polars:
+
+            >>> func(df_pd)
+               a
+            0  1
+            1  2
+            2  3
+            >>> func(df_pl).collect()
+            shape: (3, 1)
+            ┌─────┐
+            │ a   │
+            │ --- │
+            │ i64 │
+            ╞═════╡
+            │ 1   │
+            │ 2   │
+            │ 3   │
+            └─────┘
+        """
+        return super().pipe(function, *args, **kwargs)
+
     @property
     def schema(self) -> dict[str, DType]:
         r"""


Test patch:
diff --git a/tests/frame/__init__.py b/tests/frame/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/frame/pipe_test.py b/tests/frame/pipe_test.py
new file mode 100644
index 00000000..3bd6bcd7
--- /dev/null
+++ b/tests/frame/pipe_test.py
@@ -0,0 +1,29 @@
+from typing import Any
+
+import pandas as pd
+import polars as pl
+import pytest
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+data = {
+    "a": ["foo", "bars"],
+    "ab": ["foo", "bars"],
+}
+
+
+@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
+def test_str_head(constructor: Any) -> None:
+    result = nw.from_native(constructor(data)).pipe(
+        lambda _df: _df.select([x for x in _df.columns if len(x) == 2])
+    )
+    expected = {"ab": ["foo", "bars"]}
+    compare_dicts(result, expected)
+    result = (
+        nw.from_native(constructor(data))
+        .lazy()
+        .pipe(lambda _df: _df.select([x for x in _df.columns if len(x) == 2]))
+    )
+    expected = {"ab": ["foo", "bars"]}
+    compare_dicts(result, expected)