diff --git a/taskflow/tests/unit/test_types.py b/taskflow/tests/unit/test_types.py index 9399c893..d991b684 100644 --- a/taskflow/tests/unit/test_types.py +++ b/taskflow/tests/unit/test_types.py @@ -390,6 +390,18 @@ CEO root = tree.Node("josh") self.assertTrue(root.empty()) + def test_after_frozen(self): + root = tree.Node("josh") + root.add(tree.Node("josh.1")) + root.freeze() + self.assertTrue( + all(n.frozen for n in root.dfs_iter(include_self=True))) + self.assertRaises(tree.FrozenNode, + root.remove, "josh.1") + self.assertRaises(tree.FrozenNode, root.disassociate) + self.assertRaises(tree.FrozenNode, root.add, + tree.Node("josh.2")) + def test_removal(self): root = self._make_species() self.assertIsNotNone(root.remove('reptile')) diff --git a/taskflow/types/tree.py b/taskflow/types/tree.py index 56c96bbb..e0f61670 100644 --- a/taskflow/types/tree.py +++ b/taskflow/types/tree.py @@ -200,6 +200,7 @@ class Node(object): only_direct=only_direct, include_self=include_self) + @misc.disallow_when_frozen(FrozenNode) def disassociate(self): """Removes this node from its parent (if any). @@ -219,6 +220,7 @@ class Node(object): occurrences += 1 return occurrences + @misc.disallow_when_frozen(FrozenNode) def remove(self, item, only_direct=False, include_self=True): """Removes a item from this nodes children.