address review findings

Author: gonzalocasas

src/compas/datastructures/assembly/assembly.py

@@ -30,9 +30,10 @@ class Assembly(Datastructure):
 
     def __init__(self, name=None, **kwargs):
         super(Assembly, self).__init__()
-        self.attributes = {"name": name or "Assembly", "guid_key_map": {}}
+        self.attributes = {"name": name or "Assembly"}
         self.attributes.update(kwargs)
         self.graph = Graph()
+        self._parts = {}
 
     # ==========================================================================
     # data
@@ -65,6 +66,7 @@ def data(self):
     def data(self, data):
         self.attributes.update(data["attributes"] or {})
         self.graph.data = data["graph"]
+        self._parts = {part.guid: part.key for part in self.parts()}
 
     # ==========================================================================
     # properties
@@ -115,12 +117,11 @@ def add_part(self, part, key=None, **kwargs):
             The identifier of the part in the current assembly graph.
 
         """
-        guid_key_map = self.attributes["guid_key_map"]
-        if part.guid in guid_key_map.keys():
+        if part.guid in self._parts.keys():
             raise AssemblyError("Part already added to the assembly")
         key = self.graph.add_node(key=key, part=part, **kwargs)
         part.key = key
-        guid_key_map[part.guid] = part.key
+        self._parts[part.guid] = part.key
         return key
 
     def add_connection(self, a, b, **kwargs):
@@ -198,7 +199,7 @@ def find(self, guid):
             or None if the part can't be found.
 
         """
-        key = self.attributes["guid_key_map"].get(guid)
+        key = self._parts.get(guid)
 
         if key is None:
             return None

tests/compas/datastructures/test_assembly.py

@@ -67,3 +67,11 @@ def test_find_by_key():
     assert assembly.find_by_key("6") == part
 
     assert assembly.find_by_key("100") is None
+
+
+def test_find_by_key_after_deserialization():
+    assembly = Assembly()
+    part = Part()
+    assembly.add_part(part, key=2)
+    assembly = Assembly.from_data(assembly.to_data())
+    assert assembly.find_by_key(2) == part