diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index 59984c5afa3..93ad49bd223 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -1060,6 +1060,50 @@ def drop_nodes( result._replace_node(children=children_to_keep) return result + def move(self, origin: str, destination: str, parents: bool = False) -> DataTree: + """ + Move a node to a different location in the tree. + + Parameters + ---------- + origin: str + The node to move. + destination: str + The new node destination. + parents: bool, optional + If `parents` is `True`, create the additional paths needed to reach the + destination. Otherwise, raise an error. Set to `False` by default. + + Returns + ------- + DataTree + + Examples + -------- + >>> dt = DataTree.from_dict( + ... { + ... "/to_move/A": None, + ... "/to_move/B": None, + ... "/destination/other": None, + ... } + ... ) + >>> dt.move("to_move", "destination") + + Group: / + └── Group: /destination + ├── Group: /destination/other + └── Group: /destination/to_move + ├── Group: /destination/to_move/A + └── Group: /destination/to_move/B + """ + result = self.copy() + to_move = result[origin] + to_move.orphan() + result._set_item( + f"{destination}/{to_move.name}", to_move, new_nodes_along_path=parents + ) + return result + @classmethod def from_dict( cls, diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 9a15376a1f8..4ce4dd4b149 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -972,6 +972,41 @@ def test_assign(self): result = dt.assign({"foo": xr.DataArray(0), "a": DataTree()}) assert_equal(result, expected) + def test_move(self): + dt = DataTree.from_dict( + {"/to_move/A": None, "/to_move/B": None, "/destination/other": None} + ) + result = dt.move("to_move", "destination") + expected = DataTree.from_dict( + { + "/destination/to_move/A": None, + "/destination/to_move/B": None, + "/destination/other": None, + } + ) + assert_equal(result, expected) + + def test_move_no_parents(self): + dt = DataTree.from_dict( + {"/to_move/A": None, "/to_move/B": None, "/destination/other": None} + ) + with pytest.raises(KeyError, match="Could not reach node"): + dt.move("to_move", "destination/deep/down") + + def test_move_parents(self): + dt = DataTree.from_dict( + {"/to_move/A": None, "/to_move/B": None, "/destination/other": None} + ) + result = dt.move("to_move", "destination/deep/down", parents=True) + expected = DataTree.from_dict( + { + "/destination/deep/down/to_move/A": None, + "/destination/deep/down/to_move/B": None, + "/destination/other": None, + } + ) + assert_equal(result, expected) + class TestPipe: def test_noop(self, create_test_datatree):