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