Fix loop bug (#105)

* test to check for loops

* fix bug when assigning new parent

* refactor to use _is_descendant_of method

* also check children setter for loops

* whatsnew

Author: TomNicholas

datatree/tests/test_treenode.py

@@ -18,6 +18,26 @@ def test_parenting(self):
         assert mary.parent == john
         assert john.children["Mary"] is mary
 
+    def test_no_time_traveller_loops(self):
+        john = TreeNode()
+
+        with pytest.raises(TreeError, match="cannot be a parent of itself"):
+            john._set_parent(john, "John")
+
+        with pytest.raises(TreeError, match="cannot be a parent of itself"):
+            john.children = {"John": john}
+
+        mary = TreeNode()
+        rose = TreeNode()
+        mary._set_parent(john, "Mary")
+        rose._set_parent(mary, "Rose")
+
+        with pytest.raises(TreeError, match="is already a descendant"):
+            john._set_parent(rose, "John")
+
+        with pytest.raises(TreeError, match="is already a descendant"):
+            rose.children = {"John": john}
+
     def test_parent_swap(self):
         john = TreeNode()
         mary = TreeNode()

datatree/treenode.py

@@ -111,12 +111,15 @@ def _check_loop(self, new_parent: Tree | None) -> None:
                     f"Cannot set parent, as node {self} cannot be a parent of itself."
                 )
 
-            _self, *lineage = list(self.lineage)
-            if any(child is self for child in lineage):
+            if self._is_descendant_of(new_parent):
                 raise TreeError(
-                    f"Cannot set parent, as node {self} is already a descendant of node {new_parent}."
+                    f"Cannot set parent, as node {new_parent.name} is already a descendant of this node."
                 )
 
+    def _is_descendant_of(self, node: Tree) -> bool:
+        _self, *lineage = list(node.lineage)
+        return any(n is self for n in lineage)
+
     def _detach(self, parent: Tree | None) -> None:
         if parent is not None:
             self._pre_detach(parent)

docs/source/whats-new.rst

@@ -32,6 +32,9 @@ Deprecations
 Bug fixes
 ~~~~~~~~~
 
+- Fixed bug with checking that assigning parent or new children did not create a loop in the tree (:pull:`105`)
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
+
 Documentation
 ~~~~~~~~~~~~~