diff --git a/datatree/datatree.py b/datatree/datatree.py index fb26dff0..85cec7d9 100644 --- a/datatree/datatree.py +++ b/datatree/datatree.py @@ -30,7 +30,7 @@ from xarray.core.indexes import Index, Indexes from xarray.core.merge import dataset_update_method from xarray.core.options import OPTIONS as XR_OPTS -from xarray.core.utils import Default, Frozen, _default +from xarray.core.utils import Default, Frozen, _default, either_dict_or_kwargs from xarray.core.variable import Variable, calculate_dimensions from . import formatting, formatting_html @@ -821,6 +821,48 @@ def update(self, other: Dataset | Mapping[str, DataTree | DataArray]) -> None: inplace=True, children=merged_children, **vars_merge_result._asdict() ) + def assign( + self, items: Mapping[Any, Any] | None = None, **items_kwargs: Any + ) -> DataTree: + """ + Assign new data variables or child nodes to a DataTree, returning a new object + with all the original items in addition to the new ones. + + Parameters + ---------- + items : mapping of hashable to Any + Mapping from variable or child node names to the new values. If the new values + are callable, they are computed on the Dataset and assigned to new + data variables. If the values are not callable, (e.g. a DataTree, DataArray, + scalar, or array), they are simply assigned. + **items_kwargs + The keyword arguments form of ``variables``. + One of variables or variables_kwargs must be provided. + + Returns + ------- + dt : DataTree + A new DataTree with the new variables or children in addition to all the + existing items. + + Notes + ----- + Since ``kwargs`` is a dictionary, the order of your arguments may not + be preserved, and so the order of the new variables is not well-defined. + Assigning multiple items within the same ``assign`` is + possible, but you cannot reference other variables created within the + same ``assign`` call. + + See Also + -------- + xarray.Dataset.assign + pandas.DataFrame.assign + """ + items = either_dict_or_kwargs(items, items_kwargs, "assign") + dt = self.copy() + dt.update(items) + return dt + def drop_nodes( self: DataTree, names: str | Iterable[str], *, errors: ErrorOptions = "raise" ) -> DataTree: diff --git a/datatree/ops.py b/datatree/ops.py index eabc1faf..d6ac4f83 100644 --- a/datatree/ops.py +++ b/datatree/ops.py @@ -68,7 +68,6 @@ "combine_first", "reduce", "map", - "assign", "diff", "shift", "roll", diff --git a/datatree/tests/test_datatree.py b/datatree/tests/test_datatree.py index 41742085..9fa1b91f 100644 --- a/datatree/tests/test_datatree.py +++ b/datatree/tests/test_datatree.py @@ -206,6 +206,17 @@ def test_getitem_dict_like_selection_access_to_dataset(self): class TestUpdate: + def test_update(self): + dt = DataTree() + dt.update({"foo": xr.DataArray(0), "a": DataTree()}) + expected = DataTree.from_dict({"/": xr.Dataset({"foo": 0}), "a": None}) + print(dt) + print(dt.children) + print(dt._children) + print(dt["a"]) + print(expected) + dtt.assert_equal(dt, expected) + def test_update_new_named_dataarray(self): da = xr.DataArray(name="temp", data=[0, 50]) folder1 = DataTree(name="folder1") @@ -542,6 +553,18 @@ def test_drop_nodes(self): childless = dropped.drop_nodes(names=["Mary", "Ashley"], errors="ignore") assert childless.children == {} + def test_assign(self): + dt = DataTree() + expected = DataTree.from_dict({"/": xr.Dataset({"foo": 0}), "/a": None}) + + # kwargs form + result = dt.assign(foo=xr.DataArray(0), a=DataTree()) + dtt.assert_equal(result, expected) + + # dict form + result = dt.assign({"foo": xr.DataArray(0), "a": DataTree()}) + dtt.assert_equal(result, expected) + class TestPipe: def test_noop(self, create_test_datatree): diff --git a/docs/source/api.rst b/docs/source/api.rst index 3aec6d6c..23f90326 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -109,7 +109,7 @@ Manipulate the contents of all nodes in a tree simultaneously. :toctree: generated/ DataTree.copy - DataTree.assign + DataTree.assign_coords DataTree.merge DataTree.rename @@ -130,6 +130,7 @@ Manipulate the contents of a single DataTree node. .. autosummary:: :toctree: generated/ + DataTree.assign DataTree.drop_nodes Comparisons diff --git a/docs/source/whats-new.rst b/docs/source/whats-new.rst index 6dafa1e2..e57e31e4 100644 --- a/docs/source/whats-new.rst +++ b/docs/source/whats-new.rst @@ -39,6 +39,8 @@ Breaking changes By `Tom Nicholas `_. - Grafting a subtree onto another tree now leaves name of original subtree object unchanged (:issue:`116`, :pull:`172`, :pull:`178`). By `Tom Nicholas `_. +- Changed the :py:meth:`DataTree.assign` method to just work on the local node (:pull:`181`). + By `Tom Nicholas `_. Deprecations ~~~~~~~~~~~~