use guid refs

Author: tomvanmele

src/compas_model/models/elementtree.py

@@ -1,9 +1,14 @@
+from typing import TYPE_CHECKING
 from typing import Optional
+from typing import Union
 
 from compas.datastructures import Tree
 from compas.datastructures import TreeNode
 from compas_model.elements import Element
 
+if TYPE_CHECKING:
+    from compas_model.models import Model
+
 
 class ElementNode(TreeNode):
     """Class representing nodes containing elements in an element tree.
@@ -25,12 +30,12 @@ class ElementNode(TreeNode):
 
     """
 
-    tree: "ElementTree"
+    tree: "ElementTree"  # type: ignore
 
     @property
     def __data__(self) -> dict:
         data = super().__data__
-        data["element"] = None if not self.element else str(self.element.guid)
+        data["element"] = self._element
         return data
 
     @classmethod
@@ -39,16 +44,28 @@ def __from_data__(cls, data: dict) -> "ElementNode":
 
     def __init__(self, element: Optional[Element] = None, **kwargs) -> None:
         super().__init__(**kwargs)
-
-        if element:
-            element.treenode = self
+        self._element = None
         self.element = element
 
+    @property
+    def element(self):
+        if self._element:
+            return self.tree.model._elements[self._element]
+
+    @element.setter
+    def element(self, element: Optional[Union[Element, str]] = None) -> None:
+        if isinstance(element, Element):
+            self._element = str(element.guid)
+        else:
+            self._element = element
+
     def __getitem__(self, index: int) -> "ElementNode":
         return self.children[index]
 
     def __repr__(self):
         if self.parent:
+            if not self.element:
+                raise Exception("An element node that is not the root node should have an element.")
             return f"{self.element.__class__.__name__}(name={self.element.name})"
         else:
             return "ROOT"
@@ -75,28 +92,42 @@ class ElementTree(Tree):
 
     """
 
-    @property
-    def __data__(self) -> dict:
-        data = super().__data__
-        return data
+    model: "Model"
 
     @classmethod
     def __from_data__(cls, data: dict) -> "ElementTree":
         raise Exception("Serialisation outside model context not allowed.")
 
-    def __init__(self, name: Optional[str] = None) -> None:
+    def __init__(self, model: "Model", name: Optional[str] = None) -> None:
         super().__init__(name=name)
-
+        self.model = model
         root = ElementNode(name="root")
         self.add(root)
 
     @property
     def elements(self) -> list[Element]:
-        return [node.element for node in self.nodes if isinstance(node, ElementNode) if node.element]
+        return [node.element for node in self.nodes if isinstance(node, ElementNode) and node.element]
 
     @property
     def rootelements(self) -> list[Element]:
-        return [node.element for node in self.root.children if isinstance(node, ElementNode) if node.element]
+        if self.root:
+            return [node.element for node in self.root.children if isinstance(node, ElementNode) and node.element]
+        return []
+
+    def add_element(self, element: Element, parent: Optional[Union[Element, ElementNode]] = None) -> ElementNode:
+        if parent is None:
+            parentnode = self.root
+        elif isinstance(parent, Element):
+            parentnode = parent.treenode
+            if parentnode is None:
+                raise ValueError("The parent element is not part of this model.")
+        else:
+            parentnode = parent
+
+        treenode = ElementNode(element=element)
+        element.treenode = treenode
+        parentnode.add(treenode)  # type: ignore
+        return treenode
 
     def find_element_node(self, element: Element) -> ElementNode:
         """Find the node containing the element.

src/compas_model/models/interactiongraph.py

@@ -1,12 +1,10 @@
-# from typing import Generator
-from typing import Optional
+from typing import TYPE_CHECKING
 
 from compas.datastructures import Graph
 from compas_model.elements import Element  # noqa: F401
 
-# Ideally, graph (and mesh) are rewritten to use dedicated classes for nodes and edges.
-# This will allow more fine-grained control over the (types of) attributes added to nodes and edges.
-# It will also provide a much more intuitive API.
+if TYPE_CHECKING:
+    from compas_model.models import Model
 
 
 class InteractionGraph(Graph):
@@ -26,57 +24,16 @@ class InteractionGraph(Graph):
 
     """
 
-    @property
-    def __data__(self) -> dict:
-        data = super().__data__
-
-        for node, attr in data["node"].items():
-            # this modifies the attributes in place
-            # as a consequence, after accessing the __data__ property of the graph,
-            # the graph is broken
-            # to prevent this, the attribute dict should be copied
-            attr = attr.copy()
-            attr["element"] = str(attr["element"].guid)
-            data["node"][node] = attr
-        return data
-
-    @classmethod
-    def __from_data__(cls, data: dict, guid_element: dict[str, Element]) -> "InteractionGraph":
-        graph = super(InteractionGraph, cls).__from_data__(data)
-        for node, attr in graph.nodes(data=True):
-            element = guid_element[attr["element"]]
-            attr["element"] = element  # type: ignore
-            element.graphnode = node
-        return graph
-
-    def copy(self) -> "InteractionGraph":
-        # A custom implementation of copy is needed to allow passing the element dictionary to __from_data__.
-        guid_element = {}
-        for _, node in self.nodes(data=True):
-            element = node["element"]
-            guid_element[str(element.guid)] = element
-        return self.__from_data__(self.__data__, guid_element)
-
-    def __init__(
-        self,
-        default_node_attributes: Optional[dict] = None,
-        default_edge_attributes: Optional[dict] = None,
-        name: Optional[str] = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            default_node_attributes=default_node_attributes,
-            default_edge_attributes=default_edge_attributes,
-            name=name,
-            **kwargs,
-        )
+    model: "Model"
 
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
         self.update_default_node_attributes(element=None)
-        self.update_default_edge_attributes(interactions=None)
+        self.update_default_edge_attributes(modifiers=None, contacts=None)
+        self.model = None  # type: ignore
 
     def __str__(self) -> str:
         output = super().__str__()
-
         output += "\n"
         output += self._build_interactions_str()
         return output
@@ -91,23 +48,18 @@ def _build_interactions_str(self) -> str:
                     edge = nbr, node
 
                 lines.append(
-                    "- {}: {} {} {}".format(
+                    "- {}: {} {}".format(
                         nbr,
                         self.edge_attribute(edge, "modifiers"),
                         self.edge_attribute(edge, "contacts"),
-                        self.edge_attribute(edge, "collisons"),
                     )  # type: ignore
                 )
         return "\n".join(lines) + "\n"
 
-    def clear_edges(self):
-        """Clear all the edges and connectivity information of the graph."""
-        for u, v in list(self.edges()):
-            del self.edge[u][v]
-            if v in self.adjacency[u]:
-                del self.adjacency[u][v]
-            if u in self.adjacency[v]:
-                del self.adjacency[v][u]
+    def add_element(self, element: Element) -> int:
+        node = self.add_node(element=str(element.guid))
+        element.graphnode = node
+        return node
 
     def node_element(self, node: int) -> Element:
         """Get the element associated with the node.
@@ -122,4 +74,14 @@ def node_element(self, node: int) -> Element:
         :class:`compas_model.elements.Element`
 
         """
-        return self.node_attribute(node, "element")  # type: ignore
+        guid: str = self.node_attribute(node, "element")  # type: ignore
+        return self.model._elements[guid]
+
+    def clear_edges(self):
+        """Clear all the edges and connectivity information of the graph."""
+        for u, v in list(self.edges()):
+            del self.edge[u][v]
+            if v in self.adjacency[u]:
+                del self.adjacency[u][v]
+            if u in self.adjacency[v]:
+                del self.adjacency[v][u]