Allow assigning Coercible values + NamedNode internal class (#115)

* test assigning int

* allow assigning coercible values

* refactor name-related methods to intermediate class

* refactor tests to match

* fix now-exposed bug with naming

* moved test showing lack of name permanence

* whatsnew

Author: TomNicholas

datatree/datatree.py

@@ -41,7 +41,7 @@
     MappedDataWithCoords,
 )
 from .render import RenderTree
-from .treenode import NodePath, Tree, TreeNode
+from .treenode import NamedNode, NodePath, Tree
 
 if TYPE_CHECKING:
     from xarray.core.merge import CoercibleValue
@@ -228,7 +228,7 @@ def _replace(
 
 
 class DataTree(
-    TreeNode,
+    NamedNode,
     MappedDatasetMethodsMixin,
     MappedDataWithCoords,
     DataTreeArithmeticMixin,
@@ -343,20 +343,6 @@ def __init__(
         )
         self._close = ds._close
 
-    @property
-    def name(self) -> str | None:
-        """The name of this node."""
-        return self._name
-
-    @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
-
     @property
     def parent(self: DataTree) -> DataTree | None:
         """Parent of this node."""
@@ -699,14 +685,17 @@ def _set(self, key: str, val: DataTree | CoercibleValue) -> None:
         if isinstance(val, DataTree):
             val.name = key
             val.parent = self
-        elif isinstance(val, (DataArray, Variable)):
-            # TODO this should also accomodate other types that can be coerced into Variables
-            self.update({key: val})
         else:
-            raise TypeError(f"Type {type(val)} cannot be assigned to a DataTree")
+            if not isinstance(val, (DataArray, Variable)):
+                # accommodate other types that can be coerced into Variables
+                val = DataArray(val)
+
+            self.update({key: val})
 
     def __setitem__(
-        self, key: str, value: DataTree | Dataset | DataArray | Variable
+        self,
+        key: str,
+        value: Any,
     ) -> None:
         """
         Add either a child node or an array to the tree, at any position.

datatree/tests/test_datatree.py

@@ -249,13 +249,6 @@ def test_setitem_unnamed_child_node_becomes_named(self):
         john2["sonny"] = DataTree()
         assert john2["sonny"].name == "sonny"
 
-    @pytest.mark.xfail(reason="bug with name overwriting")
-    def test_setitem_child_node_keeps_name(self):
-        john = DataTree(name="john")
-        r2d2 = DataTree(name="R2D2")
-        john["Mary"] = r2d2
-        assert r2d2.name == "R2D2"
-
     def test_setitem_new_grandchild_node(self):
         john = DataTree(name="john")
         DataTree(name="mary", parent=john)
@@ -314,6 +307,17 @@ def test_setitem_unnamed_dataarray(self):
         folder1["results"] = data
         xrt.assert_equal(folder1["results"], data)
 
+    def test_setitem_variable(self):
+        var = xr.Variable(data=[0, 50], dims="x")
+        folder1 = DataTree(name="folder1")
+        folder1["results"] = var
+        xrt.assert_equal(folder1["results"], xr.DataArray(var))
+
+    def test_setitem_coerce_to_dataarray(self):
+        folder1 = DataTree(name="folder1")
+        folder1["results"] = 0
+        xrt.assert_equal(folder1["results"], xr.DataArray(0))
+
     def test_setitem_add_new_variable_to_empty_node(self):
         results = DataTree(name="results")
         results["pressure"] = xr.DataArray(data=[2, 3])

datatree/tests/test_treenode.py

@@ -1,7 +1,7 @@
 import pytest
 
 from datatree.iterators import LevelOrderIter, PreOrderIter
-from datatree.treenode import TreeError, TreeNode
+from datatree.treenode import NamedNode, TreeError, TreeNode
 
 
 class TestFamilyTree:
@@ -143,43 +143,6 @@ def test_get_from_root(self):
         assert sue._get_item("/Mary") is mary
 
 
-class TestPaths:
-    def test_path_property(self):
-        sue = TreeNode()
-        mary = TreeNode(children={"Sue": sue})
-        john = TreeNode(children={"Mary": mary})  # noqa
-        assert sue.path == "/Mary/Sue"
-        assert john.path == "/"
-
-    def test_path_roundtrip(self):
-        sue = TreeNode()
-        mary = TreeNode(children={"Sue": sue})
-        john = TreeNode(children={"Mary": mary})  # noqa
-        assert john._get_item(sue.path) == sue
-
-    def test_same_tree(self):
-        mary = TreeNode()
-        kate = TreeNode()
-        john = TreeNode(children={"Mary": mary, "Kate": kate})  # noqa
-        assert mary.same_tree(kate)
-
-    def test_relative_paths(self):
-        sue = TreeNode()
-        mary = TreeNode(children={"Sue": sue})
-        annie = TreeNode()
-        john = TreeNode(children={"Mary": mary, "Annie": annie})
-
-        assert sue.relative_to(john) == "Mary/Sue"
-        assert john.relative_to(sue) == "../.."
-        assert annie.relative_to(sue) == "../../Annie"
-        assert sue.relative_to(annie) == "../Mary/Sue"
-        assert sue.relative_to(sue) == "."
-
-        evil_kate = TreeNode()
-        with pytest.raises(ValueError, match="nodes do not lie within the same tree"):
-            sue.relative_to(evil_kate)
-
-
 class TestSetNodes:
     def test_set_child_node(self):
         john = TreeNode()
@@ -261,16 +224,66 @@ def test_del_child(self):
             del john["Mary"]
 
 
+class TestNames:
+    def test_child_gets_named_on_attach(self):
+        sue = NamedNode()
+        mary = NamedNode(children={"Sue": sue})  # noqa
+        assert sue.name == "Sue"
+
+    @pytest.mark.xfail(reason="requires refactoring to retain name")
+    def test_grafted_subtree_retains_name(self):
+        subtree = NamedNode("original")
+        root = NamedNode(children={"new_name": subtree})  # noqa
+        assert subtree.name == "original"
+
+
+class TestPaths:
+    def test_path_property(self):
+        sue = NamedNode()
+        mary = NamedNode(children={"Sue": sue})
+        john = NamedNode(children={"Mary": mary})  # noqa
+        assert sue.path == "/Mary/Sue"
+        assert john.path == "/"
+
+    def test_path_roundtrip(self):
+        sue = NamedNode()
+        mary = NamedNode(children={"Sue": sue})
+        john = NamedNode(children={"Mary": mary})  # noqa
+        assert john._get_item(sue.path) == sue
+
+    def test_same_tree(self):
+        mary = NamedNode()
+        kate = NamedNode()
+        john = NamedNode(children={"Mary": mary, "Kate": kate})  # noqa
+        assert mary.same_tree(kate)
+
+    def test_relative_paths(self):
+        sue = NamedNode()
+        mary = NamedNode(children={"Sue": sue})
+        annie = NamedNode()
+        john = NamedNode(children={"Mary": mary, "Annie": annie})
+
+        assert sue.relative_to(john) == "Mary/Sue"
+        assert john.relative_to(sue) == "../.."
+        assert annie.relative_to(sue) == "../../Annie"
+        assert sue.relative_to(annie) == "../Mary/Sue"
+        assert sue.relative_to(sue) == "."
+
+        evil_kate = NamedNode()
+        with pytest.raises(ValueError, match="nodes do not lie within the same tree"):
+            sue.relative_to(evil_kate)
+
+
 def create_test_tree():
-    f = TreeNode()
-    b = TreeNode()
-    a = TreeNode()
-    d = TreeNode()
-    c = TreeNode()
-    e = TreeNode()
-    g = TreeNode()
-    i = TreeNode()
-    h = TreeNode()
+    f = NamedNode()
+    b = NamedNode()
+    a = NamedNode()
+    d = NamedNode()
+    c = NamedNode()
+    e = NamedNode()
+    g = NamedNode()
+    i = NamedNode()
+    h = NamedNode()
 
     f.children = {"b": b, "g": g}
     b.children = {"a": a, "d": d}
@@ -286,7 +299,7 @@ def test_preorderiter(self):
         tree = create_test_tree()
         result = [node.name for node in PreOrderIter(tree)]
         expected = [
-            None,  # root TreeNode is unnamed
+            None,  # root Node is unnamed
             "b",
             "a",
             "d",
@@ -302,7 +315,7 @@ def test_levelorderiter(self):
         tree = create_test_tree()
         result = [node.name for node in LevelOrderIter(tree)]
         expected = [
-            None,  # root TreeNode is unnamed
+            None,  # root Node is unnamed
             "b",
             "g",
             "a",
@@ -317,19 +330,19 @@ def test_levelorderiter(self):
 
 class TestRenderTree:
     def test_render_nodetree(self):
-        sam = TreeNode()
-        ben = TreeNode()
-        mary = TreeNode(children={"Sam": sam, "Ben": ben})
-        kate = TreeNode()
-        john = TreeNode(children={"Mary": mary, "Kate": kate})
+        sam = NamedNode()
+        ben = NamedNode()
+        mary = NamedNode(children={"Sam": sam, "Ben": ben})
+        kate = NamedNode()
+        john = NamedNode(children={"Mary": mary, "Kate": kate})
 
         printout = john.__str__()
         expected_nodes = [
-            "TreeNode()",
-            "TreeNode('Mary')",
-            "TreeNode('Sam')",
-            "TreeNode('Ben')",
-            "TreeNode('Kate')",
+            "NamedNode()",
+            "NamedNode('Mary')",
+            "NamedNode('Sam')",
+            "NamedNode('Ben')",
+            "NamedNode('Kate')",
         ]
         for expected_node, printed_node in zip(expected_nodes, printout.splitlines()):
             assert expected_node in printed_node

datatree/treenode.py

@@ -113,7 +113,7 @@ def _check_loop(self, new_parent: Tree | None) -> None:
 
             if self._is_descendant_of(new_parent):
                 raise TreeError(
-                    f"Cannot set parent, as node {new_parent.name} is already a descendant of this node."
+                    "Cannot set parent, as intended parent is already a descendant of this node."
                 )
 
     def _is_descendant_of(self, node: Tree) -> bool:
@@ -461,18 +461,72 @@ def update(self: Tree, other: Mapping[str, Tree]) -> None:
         new_children = {**self.children, **other}
         self.children = new_children
 
+    def same_tree(self, other: Tree) -> bool:
+        """True if other node is in the same tree as this node."""
+        return self.root is other.root
+
+    def find_common_ancestor(self, other: Tree) -> Tree:
+        """
+        Find the first common ancestor of two nodes in the same tree.
+
+        Raise ValueError if they are not in the same tree.
+        """
+        common_ancestor = None
+        for node in other.iter_lineage():
+            if node in self.ancestors:
+                common_ancestor = node
+                break
+
+        if not common_ancestor:
+            raise ValueError(
+                "Cannot find relative path because nodes do not lie within the same tree"
+            )
+
+        return common_ancestor
+
+    def _path_to_ancestor(self, ancestor: Tree) -> NodePath:
+        generation_gap = list(self.lineage).index(ancestor)
+        path_upwards = "../" * generation_gap if generation_gap > 0 else "/"
+        return NodePath(path_upwards)
+
+
+class NamedNode(TreeNode, Generic[Tree]):
+    """
+    A TreeNode which knows its own name.
+
+    Implements path-like relationships to other nodes in its tree.
+    """
+
+    _name: Optional[str]
+    _parent: Optional[Tree]
+    _children: OrderedDict[str, Tree]
+
+    def __init__(self, name=None, children=None):
+        super().__init__(children=children)
+        self._name = None
+        self.name = name
+
     @property
     def name(self) -> str | None:
-        """If node has a parent, this is the key under which it is stored in `parent.children`."""
-        if self.parent:
-            return next(
-                name for name, child in self.parent.children.items() if child is self
-            )
-        else:
-            return None
+        """The name of this node."""
+        return self._name
+
+    @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
 
     def __str__(self) -> str:
-        return f"TreeNode({self.name})" if self.name else "TreeNode()"
+        return f"NamedNode({self.name})" if self.name else "NamedNode()"
+
+    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
 
     @property
     def path(self) -> str:
@@ -483,9 +537,9 @@ def path(self) -> str:
             root, *ancestors = self.ancestors
             # don't include name of root because (a) root might not have a name & (b) we want path relative to root.
             names = [node.name for node in ancestors]
-            return "/" + "/".join(names)  # type: ignore
+            return "/" + "/".join(names)
 
-    def relative_to(self, other: Tree) -> str:
+    def relative_to(self: NamedNode, other: NamedNode) -> str:
         """
         Compute the relative path from this node to node `other`.
 
@@ -505,31 +559,3 @@ def relative_to(self, other: Tree) -> str:
             return str(
                 path_to_common_ancestor / this_path.relative_to(common_ancestor.path)
             )
-
-    def same_tree(self, other: Tree) -> bool:
-        """True if other node is in the same tree as this node."""
-        return self.root is other.root
-
-    def find_common_ancestor(self, other: Tree) -> Tree:
-        """
-        Find the first common ancestor of two nodes in the same tree.
-
-        Raise ValueError if they are not in the same tree.
-        """
-        common_ancestor = None
-        for node in other.iter_lineage():
-            if node in self.ancestors:
-                common_ancestor = node
-                break
-
-        if not common_ancestor:
-            raise ValueError(
-                "Cannot find relative path because nodes do not lie within the same tree"
-            )
-
-        return common_ancestor
-
-    def _path_to_ancestor(self, ancestor: Tree) -> NodePath:
-        generation_gap = list(self.lineage).index(ancestor)
-        path_upwards = "../" * generation_gap if generation_gap > 0 else "/"
-        return NodePath(path_upwards)

docs/source/whats-new.rst

@@ -42,6 +42,8 @@ Bug fixes
 - Modifying the contents of a ``DataTree`` object via the ``DataTree.ds`` attribute is now forbidden, which prevents
   any possibility of the contents of a ``DataTree`` object and its ``.ds`` attribute diverging. (:issue:`38`, :pull:`99`)
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
+- Fixed a bug so that names of children now always match keys under which parents store them (:pull:`99`).
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
 
 Documentation
 ~~~~~~~~~~~~~
@@ -56,10 +58,14 @@ Internal Changes
   This approach means that the ``DataTree`` class now effectively copies and extends the internal structure of
   ``xarray.Dataset``. (:pull:`41`)
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
+- Refactored to use intermediate ``NamedNode`` class, separating implementation of methods requiring a ``name``
+  attribute from those not requiring it.
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
 - Made ``testing.test_datatree.create_test_datatree`` into a pytest fixture (:pull:`107`).
   By `Benjamin Woods <https://github.com/benjaminwoods>`_.
 
 
+
 .. _whats-new.v0.0.6:
 
 v0.0.6 (06/03/2022)