Shallow copy parent and children in DataTree constructor (pydata#9297)

* add tests

* fix by shallow copying

* correct first few tests

* replace constructors in tests with DataTree.from_dict

* rewrite simple_datatree fixture to use DataTree.from_dict

* fix incorrect creation of nested tree in formatting test

* Update doctests for from_dict constructor

* swap i and h in doctest example for clarity.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix a few mypy errors.

* Bonkers way to set type checking

I will happily take something better.

But this was the error I was getting

xarray/tests/test_datatree.py:127: error: Argument 1 to "relative_to" of
"NamedNode" has incompatible type "DataTree[Any] | DataArray"; expected
"NamedNode[Any]"  [arg-type]

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Removes parent keyword from DataTree constructor

But it doesn't fix all the tests There's three tests that I don't fully know
what should be tested or if they still make sense.

* fix test_setparent_unnamed_child_node_fails

* fix test_dont_modify_parent_inplace -> bug?

* fix test_create_two_children

* make .parent read-only, and remove tests which test the parent setter

* update error message to reflect fact that .children is Frozen

* fix another test

* add test that parent setter tells you to set children instead

* fix mypy error due to overriding settable property with read-only property

* fix test by not trying to set parent via kwarg

---------

Co-authored-by: Matt Savoie <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

xarray/core/datatree.py

@@ -423,7 +423,6 @@ class DataTree(
     def __init__(
         self,
         data: Dataset | DataArray | None = None,
-        parent: DataTree | None = None,
         children: Mapping[str, DataTree] | None = None,
         name: str | None = None,
     ):
@@ -439,8 +438,6 @@ def __init__(
         data : Dataset, DataArray, or None, optional
             Data to store under the .ds attribute of this node. DataArrays will
             be promoted to Datasets. Default is None.
-        parent : DataTree, optional
-            Parent node to this node. Default is None.
         children : Mapping[str, DataTree], optional
             Any child nodes of this node. Default is None.
         name : str, optional
@@ -459,8 +456,9 @@ def __init__(
 
         super().__init__(name=name)
         self._set_node_data(_coerce_to_dataset(data))
-        self.parent = parent
-        self.children = children
+
+        # shallow copy to avoid modifying arguments in-place (see GH issue #9196)
+        self.children = {name: child.copy() for name, child in children.items()}
 
     def _set_node_data(self, ds: Dataset):
         data_vars, coord_vars = _collect_data_and_coord_variables(ds)
@@ -497,17 +495,6 @@ def _dims(self) -> ChainMap[Hashable, int]:
     def _indexes(self) -> ChainMap[Hashable, Index]:
         return ChainMap(self._node_indexes, *(p._node_indexes for p in self.parents))
 
-    @property
-    def parent(self: DataTree) -> DataTree | None:
-        """Parent of this node."""
-        return self._parent
-
-    @parent.setter
-    def parent(self: DataTree, new_parent: DataTree) -> None:
-        if new_parent and self.name is None:
-            raise ValueError("Cannot set an unnamed node as a child of another node")
-        self._set_parent(new_parent, self.name)
-
     def _to_dataset_view(self, rebuild_dims: bool) -> DatasetView:
         variables = dict(self._data_variables)
         variables |= self._coord_variables
@@ -896,7 +883,7 @@ def _set(self, key: str, val: DataTree | CoercibleValue) -> None:
             # create and assign a shallow copy here so as not to alter original name of node in grafted tree
             new_node = val.copy(deep=False)
             new_node.name = key
-            new_node.parent = self
+            new_node._set_parent(new_parent=self, child_name=key)
         else:
             if not isinstance(val, DataArray | Variable):
                 # accommodate other types that can be coerced into Variables
@@ -1097,7 +1084,7 @@ def from_dict(
             obj = root_data.copy()
             obj.orphan()
         else:
-            obj = cls(name=name, data=root_data, parent=None, children=None)
+            obj = cls(name=name, data=root_data, children=None)
 
         def depth(item) -> int:
             pathstr, _ = item

xarray/core/datatree_render.py

@@ -51,11 +51,16 @@ def __init__(self):
 
         >>> from xarray.core.datatree import DataTree
         >>> from xarray.core.datatree_render import RenderDataTree
-        >>> root = DataTree(name="root")
-        >>> s0 = DataTree(name="sub0", parent=root)
-        >>> s0b = DataTree(name="sub0B", parent=s0)
-        >>> s0a = DataTree(name="sub0A", parent=s0)
-        >>> s1 = DataTree(name="sub1", parent=root)
+        >>> root = DataTree.from_dict(
+        ...     {
+        ...         "/": None,
+        ...         "/sub0": None,
+        ...         "/sub0/sub0B": None,
+        ...         "/sub0/sub0A": None,
+        ...         "/sub1": None,
+        ...     },
+        ...     name="root",
+        ... )
         >>> print(RenderDataTree(root))
         <xarray.DataTree 'root'>
         Group: /
@@ -98,11 +103,16 @@ def __init__(
         >>> from xarray import Dataset
         >>> from xarray.core.datatree import DataTree
         >>> from xarray.core.datatree_render import RenderDataTree
-        >>> root = DataTree(name="root", data=Dataset({"a": 0, "b": 1}))
-        >>> s0 = DataTree(name="sub0", parent=root, data=Dataset({"c": 2, "d": 3}))
-        >>> s0b = DataTree(name="sub0B", parent=s0, data=Dataset({"e": 4}))
-        >>> s0a = DataTree(name="sub0A", parent=s0, data=Dataset({"f": 5, "g": 6}))
-        >>> s1 = DataTree(name="sub1", parent=root, data=Dataset({"h": 7}))
+        >>> root = DataTree.from_dict(
+        ...     {
+        ...         "/": Dataset({"a": 0, "b": 1}),
+        ...         "/sub0": Dataset({"c": 2, "d": 3}),
+        ...         "/sub0/sub0B": Dataset({"e": 4}),
+        ...         "/sub0/sub0A": Dataset({"f": 5, "g": 6}),
+        ...         "/sub1": Dataset({"h": 7}),
+        ...     },
+        ...     name="root",
+        ... )
 
         # Simple one line:
 
@@ -208,17 +218,16 @@ def by_attr(self, attrname: str = "name") -> str:
         >>> from xarray import Dataset
         >>> from xarray.core.datatree import DataTree
         >>> from xarray.core.datatree_render import RenderDataTree
-        >>> root = DataTree(name="root")
-        >>> s0 = DataTree(name="sub0", parent=root)
-        >>> s0b = DataTree(
-        ...     name="sub0B", parent=s0, data=Dataset({"foo": 4, "bar": 109})
+        >>> root = DataTree.from_dict(
+        ...     {
+        ...         "/sub0/sub0B": Dataset({"foo": 4, "bar": 109}),
+        ...         "/sub0/sub0A": None,
+        ...         "/sub1/sub1A": None,
+        ...         "/sub1/sub1B": Dataset({"bar": 8}),
+        ...         "/sub1/sub1C/sub1Ca": None,
+        ...     },
+        ...     name="root",
         ... )
-        >>> s0a = DataTree(name="sub0A", parent=s0)
-        >>> s1 = DataTree(name="sub1", parent=root)
-        >>> s1a = DataTree(name="sub1A", parent=s1)
-        >>> s1b = DataTree(name="sub1B", parent=s1, data=Dataset({"bar": 8}))
-        >>> s1c = DataTree(name="sub1C", parent=s1)
-        >>> s1ca = DataTree(name="sub1Ca", parent=s1c)
         >>> print(RenderDataTree(root).by_attr("name"))
         root
         ├── sub0

xarray/core/iterators.py

@@ -28,15 +28,9 @@ class LevelOrderIter(Iterator):
     --------
     >>> from xarray.core.datatree import DataTree
     >>> from xarray.core.iterators import LevelOrderIter
-    >>> f = DataTree(name="f")
-    >>> b = DataTree(name="b", parent=f)
-    >>> a = DataTree(name="a", parent=b)
-    >>> d = DataTree(name="d", parent=b)
-    >>> c = DataTree(name="c", parent=d)
-    >>> e = DataTree(name="e", parent=d)
-    >>> g = DataTree(name="g", parent=f)
-    >>> i = DataTree(name="i", parent=g)
-    >>> h = DataTree(name="h", parent=i)
+    >>> f = DataTree.from_dict(
+    ...     {"/b/a": None, "/b/d/c": None, "/b/d/e": None, "/g/h/i": None}, name="f"
+    ... )
     >>> print(f)
     <xarray.DataTree 'f'>
     Group: /
@@ -46,19 +40,19 @@ class LevelOrderIter(Iterator):
     │       ├── Group: /b/d/c
     │       └── Group: /b/d/e
     └── Group: /g
-        └── Group: /g/i
-            └── Group: /g/i/h
+        └── Group: /g/h
+            └── Group: /g/h/i
     >>> [node.name for node in LevelOrderIter(f)]
-    ['f', 'b', 'g', 'a', 'd', 'i', 'c', 'e', 'h']
+    ['f', 'b', 'g', 'a', 'd', 'h', 'c', 'e', 'i']
     >>> [node.name for node in LevelOrderIter(f, maxlevel=3)]
-    ['f', 'b', 'g', 'a', 'd', 'i']
+    ['f', 'b', 'g', 'a', 'd', 'h']
     >>> [
     ...     node.name
     ...     for node in LevelOrderIter(f, filter_=lambda n: n.name not in ("e", "g"))
     ... ]
-    ['f', 'b', 'a', 'd', 'i', 'c', 'h']
+    ['f', 'b', 'a', 'd', 'h', 'c', 'i']
     >>> [node.name for node in LevelOrderIter(f, stop=lambda n: n.name == "d")]
-    ['f', 'b', 'g', 'a', 'i', 'h']
+    ['f', 'b', 'g', 'a', 'h', 'i']
     """
 
     def __init__(

xarray/core/treenode.py

@@ -86,6 +86,12 @@ def parent(self) -> Tree | None:
         """Parent of this node."""
         return self._parent
 
+    @parent.setter
+    def parent(self: Tree, new_parent: Tree) -> None:
+        raise AttributeError(
+            "Cannot set parent attribute directly, you must modify the children of the other node instead using dict-like syntax"
+        )
+
     def _set_parent(
         self, new_parent: Tree | None, child_name: str | None = None
     ) -> None:

xarray/tests/conftest.py

@@ -196,14 +196,17 @@ def _create_test_datatree(modify=lambda ds: ds):
         set2_data = modify(xr.Dataset({"a": ("x", [2, 3]), "b": ("x", [0.1, 0.2])}))
         root_data = modify(xr.Dataset({"a": ("y", [6, 7, 8]), "set0": ("x", [9, 10])}))
 
-        # Avoid using __init__ so we can independently test it
-        root: DataTree = DataTree(data=root_data)
-        set1: DataTree = DataTree(name="set1", parent=root, data=set1_data)
-        DataTree(name="set1", parent=set1)
-        DataTree(name="set2", parent=set1)
-        set2: DataTree = DataTree(name="set2", parent=root, data=set2_data)
-        DataTree(name="set1", parent=set2)
-        DataTree(name="set3", parent=root)
+        root = DataTree.from_dict(
+            {
+                "/": root_data,
+                "/set1": set1_data,
+                "/set1/set1": None,
+                "/set1/set2": None,
+                "/set2": set2_data,
+                "/set2/set1": None,
+                "/set3": None,
+            }
+        )
 
         return root