From b69e9b5ab79cd1f71ce58030ff56c9cb8359e630 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 3 Oct 2022 16:48:56 +0200 Subject: [PATCH 1/5] add a minimal pipe implementation --- datatree/datatree.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/datatree/datatree.py b/datatree/datatree.py index 6f78d8c8..1517ede6 100644 --- a/datatree/datatree.py +++ b/datatree/datatree.py @@ -1014,6 +1014,9 @@ def map_over_subtree_inplace( if node.has_data: node.ds = func(node.ds, *args, **kwargs) + def pipe(self, func, *args, **kwargs): + return func(self, *args, **kwargs) + def render(self): """Print tree structure, including any data stored at each node.""" for pre, fill, node in RenderTree(self): From 61a60c7c74fee5c952863de94529c689683f12fd Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 5 Oct 2022 14:07:25 +0200 Subject: [PATCH 2/5] copy the code of Dataset.pipe --- datatree/datatree.py | 61 ++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 59 insertions(+), 2 deletions(-) diff --git a/datatree/datatree.py b/datatree/datatree.py index 1517ede6..ff7a417b 100644 --- a/datatree/datatree.py +++ b/datatree/datatree.py @@ -1014,8 +1014,65 @@ def map_over_subtree_inplace( if node.has_data: node.ds = func(node.ds, *args, **kwargs) - def pipe(self, func, *args, **kwargs): - return func(self, *args, **kwargs) + def pipe( + self, func: Callable | tuple[Callable, str], *args: Any, **kwargs: Any + ) -> Any: + """Apply ``func(self, *args, **kwargs)`` + + This method replicates the pandas method of the same name. + + Parameters + ---------- + func : callable + function to apply to this xarray object (Dataset/DataArray). + ``args``, and ``kwargs`` are passed into ``func``. + Alternatively a ``(callable, data_keyword)`` tuple where + ``data_keyword`` is a string indicating the keyword of + ``callable`` that expects the xarray object. + *args + positional arguments passed into ``func``. + **kwargs + a dictionary of keyword arguments passed into ``func``. + + Returns + ------- + object : Any + the return type of ``func``. + + Notes + ----- + Use ``.pipe`` when chaining together functions that expect + xarray or pandas objects, e.g., instead of writing + + .. code:: python + + f(g(h(dt), arg1=a), arg2=b, arg3=c) + + You can write + + .. code:: python + + (dt.pipe(h).pipe(g, arg1=a).pipe(f, arg2=b, arg3=c)) + + If you have a function that takes the data as (say) the second + argument, pass a tuple indicating which keyword expects the + data. For example, suppose ``f`` takes its data as ``arg2``: + + .. code:: python + + (dt.pipe(h).pipe(g, arg1=a).pipe((f, "arg2"), arg1=a, arg3=c)) + + """ + if isinstance(func, tuple): + func, target = func + if target in kwargs: + raise ValueError( + f"{target} is both the pipe target and a keyword argument" + ) + kwargs[target] = self + else: + args = (self,) + args + return func(*args, **kwargs) def render(self): """Print tree structure, including any data stored at each node.""" From 7e08c02d509b958afacd3a1e1564c9354689351c Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 2 Nov 2022 18:39:23 +0100 Subject: [PATCH 3/5] create a documentation page --- docs/source/api.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/api.rst b/docs/source/api.rst index 209d4ab9..49caaea8 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -96,6 +96,7 @@ For manipulating, traversing, navigating, or mapping over the tree structure. DataTree.iter_lineage DataTree.find_common_ancestor map_over_subtree + DataTree.pipe DataTree Contents ----------------- From 9ddb4eef819c0c4437e3655b4228fc4407dba135 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 2 Nov 2022 18:56:18 +0100 Subject: [PATCH 4/5] add tests for `pipe` --- datatree/tests/test_datatree.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/datatree/tests/test_datatree.py b/datatree/tests/test_datatree.py index ca9fae5f..dd08618d 100644 --- a/datatree/tests/test_datatree.py +++ b/datatree/tests/test_datatree.py @@ -448,3 +448,35 @@ def test_arithmetic(self, create_test_datatree): class TestRestructuring: ... + + +class TestPipe: + def test_noop(self, create_test_datatree): + dt = create_test_datatree() + + actual = dt.pipe(lambda tree: tree) + assert actual.identical(dt) + + def test_params(self, create_test_datatree): + dt = create_test_datatree() + + def f(tree, **attrs): + return tree.assign(arr_with_attrs=xr.Variable("dim0", [], attrs=attrs)) + + attrs = {"x": 1, "y": 2, "z": 3} + + actual = dt.pipe(f, **attrs) + assert actual["arr_with_attrs"].attrs == attrs + + def test_named_self(self, create_test_datatree): + dt = create_test_datatree() + + def f(x, tree, y): + tree.attrs.update({"x": x, "y": y}) + return tree + + attrs = {"x": 1, "y": 2} + + actual = dt.pipe((f, "tree"), **attrs) + + assert actual is dt and actual.attrs == attrs From 5e8f2400efec95513f869068e3439da59c71cf33 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 2 Nov 2022 18:57:32 +0100 Subject: [PATCH 5/5] whats-new.rst --- docs/source/whats-new.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/source/whats-new.rst b/docs/source/whats-new.rst index 25bb8614..daee3fd8 100644 --- a/docs/source/whats-new.rst +++ b/docs/source/whats-new.rst @@ -25,6 +25,8 @@ New Features - Add the ability to register accessors on ``DataTree`` objects, by using ``register_datatree_accessor``. (:pull:`144`) By `Tom Nicholas `_. +- Allow method chaining with a new :py:meth:`DataTree.pipe` method (:issue:`151`, :pull:`156`). + By `Justus Magin `_. Breaking changes ~~~~~~~~~~~~~~~~