From 9928394a0c802ab810820456614131f5dca53958 Mon Sep 17 00:00:00 2001
From: Li Chen <leochen1990@gmail.com>
Date: Fri, 6 Oct 2023 12:21:38 +0200
Subject: [PATCH] make tree link dynamic

---
 src/compas/datastructures/tree/tree.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/src/compas/datastructures/tree/tree.py b/src/compas/datastructures/tree/tree.py
index 1c0a67abc6c..b99ecd47baf 100644
--- a/src/compas/datastructures/tree/tree.py
+++ b/src/compas/datastructures/tree/tree.py
@@ -26,6 +26,8 @@ class TreeNode(Data):
         The parent node of the tree node.
     children : set[:class:`~compas.datastructures.TreeNode`]
         The children of the tree node.
+    tree : :class:`~compas.datastructures.Tree`
+        The tree to which the node belongs.
     is_root : bool
         True if the node is the root node of the tree.
     is_leaf : bool
@@ -55,6 +57,7 @@ def __init__(self, name=None, attributes=None):
         self.attributes = attributes or {}
         self._parent = None
         self._children = set()
+        self._tree = None
 
     def __repr__(self):
         return "<TreeNode {}>".format(self.name)
@@ -79,6 +82,13 @@ def parent(self):
     def children(self):
         return self._children
 
+    @property
+    def tree(self):
+        if self.is_root:
+            return self._tree
+        else:
+            return self.parent.tree
+
     @property
     def data(self):
         return {
@@ -303,6 +313,7 @@ def add(self, node, parent=None):
                 raise ValueError("The tree already has a root node, remove it first.")
 
             self._root = node
+            node._tree = self
 
         else:
             # add the node as a child of the parent node
@@ -338,6 +349,7 @@ def remove(self, node):
         """
         if node == self.root:
             self._root = None
+            node._tree = None
         else:
             node.parent.remove(node)