Add assign method (#181)

* WIP assign method

* remove rogue prints

* remove assign from mapped methods

* moved assign in docs to reflect it only acting on single nodes now

* whatsnew

Author: TomNicholas

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:

datatree/ops.py

@@ -68,7 +68,6 @@
     "combine_first",
     "reduce",
     "map",
-    "assign",
     "diff",
     "shift",
     "roll",

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):

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

docs/source/whats-new.rst

@@ -39,6 +39,8 @@ Breaking changes
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
 - Grafting a subtree onto another tree now leaves name of original subtree object unchanged (:issue:`116`, :pull:`172`, :pull:`178`).
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
+- Changed the :py:meth:`DataTree.assign` method to just work on the local node (:pull:`181`).
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
 
 Deprecations
 ~~~~~~~~~~~~