diff --git a/datatree/datatree.py b/datatree/datatree.py index c86c2e2e..fa84e03e 100644 --- a/datatree/datatree.py +++ b/datatree/datatree.py @@ -1482,7 +1482,7 @@ def to_netcdf( Note that unlimited_dims may also be set via ``dataset.encoding["unlimited_dims"]``. kwargs : - Addional keyword arguments to be passed to ``xarray.Dataset.to_netcdf`` + Additional keyword arguments to be passed to ``xarray.Dataset.to_netcdf`` """ from .io import _datatree_to_netcdf diff --git a/datatree/tests/test_datatree.py b/datatree/tests/test_datatree.py index e9f373d7..5854216b 100644 --- a/datatree/tests/test_datatree.py +++ b/datatree/tests/test_datatree.py @@ -68,6 +68,57 @@ def test_child_gets_named_on_attach(self): mary = DataTree(children={"Sue": sue}) # noqa assert sue.name == "Sue" + def test_setting_node_name_keeps_tree_linkage(self): + root = DataTree(name="root") + child = DataTree(name="child", parent=root) + grandchild = DataTree(name="grandchild", parent=child) + + assert root.name == "root" + assert child.name == "child" + assert grandchild.name == "grandchild" + + # changing the name of a child node should correctly update the dict key in + # its parent's children + child.name = "childish" + + assert child.name == "childish" + assert "childish" in root + assert list(root.children) == ["childish"] + + # changing name of root + root.name = "ginger" + + assert root.name == "ginger" + + def test_setting_node_name_keeps_tree_linkage_and_order(self): + root = DataTree.from_dict( + { + "/one": None, + "/two": None, + "/three": None, + } + ) + assert list(root.children) == ["one", "two", "three"] + + second_child = root["/two"] + second_child.name = "second" + + assert second_child.name == "second" + assert "second" in root + + # Order was preserved + assert list(root.children) == ["one", "second", "three"] + + def test_setting_node_to_none(self): + root = DataTree(name="root") + child = DataTree(name="child", parent=root) + + with pytest.raises( + ValueError, + match=r"a node with a parent cannot have its name set to None", + ): + child.name = None + class TestPaths: def test_path_property(self): diff --git a/datatree/tests/test_formatting.py b/datatree/tests/test_formatting.py index 0f64644c..76c6e072 100644 --- a/datatree/tests/test_formatting.py +++ b/datatree/tests/test_formatting.py @@ -108,13 +108,13 @@ def test_diff_node_data(self): Data in nodes at position '/a' do not match: Data variables only on the left object: - v int64 1 + v int64 8B 1 Data in nodes at position '/a/b' do not match: Differing data variables: - L w int64 5 - R w int64 6""" + L w int64 8B 5 + R w int64 8B 6""" ) actual = diff_tree_repr(dt_1, dt_2, "equals") assert actual == expected diff --git a/datatree/treenode.py b/datatree/treenode.py index 1689d261..edbb408a 100644 --- a/datatree/treenode.py +++ b/datatree/treenode.py @@ -5,6 +5,7 @@ from pathlib import PurePosixPath from typing import ( TYPE_CHECKING, + Any, Generic, Iterator, Mapping, @@ -91,7 +92,7 @@ def parent(self) -> Tree | None: return self._parent def _set_parent( - self, new_parent: Tree | None, child_name: Optional[str] = None + self, new_parent: Tree | None, child_name: str | None = None ) -> None: # TODO is it possible to refactor in a way that removes this private method? @@ -596,12 +597,22 @@ def name(self) -> str | None: @name.setter def name(self, name: str | None) -> None: - if name is not None: - if not isinstance(name, str): - raise TypeError("node name must be a string or None") - if "/" in name: - raise ValueError("node names cannot contain forward slashes") - self._name = name + self._validate_name(name) + + if self._parent is None: + self._name = name + return + + if name is None: + raise ValueError("a node with a parent cannot have its name set to None") + + # Rename node while preserving the order of its parent's children + self._parent.children = OrderedDict( + ( + (k if k != self._name else name, v) + for k, v in self._parent.children.items() + ) + ) def __str__(self) -> str: return f"NamedNode({self.name})" if self.name else "NamedNode()" @@ -609,7 +620,7 @@ def __str__(self) -> str: def _post_attach(self: NamedNode, parent: NamedNode) -> None: """Ensures child has name attribute corresponding to key under which it has been stored.""" key = next(k for k, v in parent.children.items() if v is self) - self.name = key + self._name = key @property def path(self) -> str: @@ -677,3 +688,10 @@ def _path_to_ancestor(self, ancestor: NamedNode) -> NodePath: generation_gap = list(parents_paths).index(ancestor.path) path_upwards = "../" * generation_gap if generation_gap > 0 else "." return NodePath(path_upwards) + + @staticmethod + def _validate_name(name: Any): + if name is not None and not isinstance(name, str): + raise TypeError("node name must be a string or None") + if name is not None and "/" in name: + raise ValueError("node names cannot contain forward slashes") diff --git a/docs/source/whats-new.rst b/docs/source/whats-new.rst index 2f6e4f88..62f324cb 100644 --- a/docs/source/whats-new.rst +++ b/docs/source/whats-new.rst @@ -45,6 +45,10 @@ Deprecations Bug fixes ~~~~~~~~~ + +- Renaming a node also updates the key of its parent's children + (:issue:`309`, :pull:`310`) + By `Etienne Schalk `_. - Keep attributes on nodes containing no data in :py:func:`map_over_subtree`. (:issue:`278`, :pull:`279`) By `Sam Levang `_.