Skip to content

Shallow copy parent and children in DataTree constructor #9297

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
4b7e367
add tests
TomNicholas Jul 31, 2024
6c56b12
fix by shallow copying
TomNicholas Jul 31, 2024
5db25bd
correct first few tests
TomNicholas Aug 2, 2024
b28cb0c
replace constructors in tests with DataTree.from_dict
TomNicholas Aug 20, 2024
bc543b1
rewrite simple_datatree fixture to use DataTree.from_dict
TomNicholas Aug 20, 2024
a080689
fix incorrect creation of nested tree in formatting test
TomNicholas Aug 20, 2024
b73906e
Merge branch 'main' into datatree_init_dont_modify_inplace
TomNicholas Aug 20, 2024
2a6b88d
Update doctests for from_dict constructor
flamingbear Sep 2, 2024
5ddcb44
swap i and h in doctest example for clarity.
flamingbear Sep 2, 2024
6c496b9
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 2, 2024
07d6904
Fix a few mypy errors.
flamingbear Sep 2, 2024
015d2ce
Bonkers way to set type checking
flamingbear Sep 2, 2024
847f238
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 2, 2024
7c530f5
Removes parent keyword from DataTree constructor
flamingbear Sep 3, 2024
9a004d4
fix test_setparent_unnamed_child_node_fails
flamingbear Sep 3, 2024
fe5ae9c
fix test_dont_modify_parent_inplace -> bug?
flamingbear Sep 3, 2024
897b84d
Merge branch 'main' into datatree_init_dont_modify_inplace
TomNicholas Sep 7, 2024
7ce6b56
fix test_create_two_children
TomNicholas Sep 7, 2024
897b589
make .parent read-only, and remove tests which test the parent setter
TomNicholas Sep 7, 2024
9938f3b
update error message to reflect fact that .children is Frozen
TomNicholas Sep 7, 2024
cb7fdfa
fix another test
TomNicholas Sep 7, 2024
1f651cb
add test that parent setter tells you to set children instead
TomNicholas Sep 7, 2024
80b4bcd
fix mypy error due to overriding settable property with read-only pro…
TomNicholas Sep 7, 2024
f4c7243
fix test by not trying to set parent via kwarg
TomNicholas Sep 7, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 5 additions & 18 deletions xarray/core/datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
):
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
49 changes: 29 additions & 20 deletions xarray/core/datatree_render.py
Original file line number Diff line number Diff line change
Expand Up @@ -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: /
Expand Down Expand Up @@ -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:

Expand Down Expand Up @@ -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
Expand Down
24 changes: 9 additions & 15 deletions xarray/core/iterators.py
Original file line number Diff line number Diff line change
Expand Up @@ -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: /
Expand All @@ -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__(
Expand Down
6 changes: 6 additions & 0 deletions xarray/core/treenode.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
19 changes: 11 additions & 8 deletions xarray/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Loading
Loading