Merge pull request #1032 from compas-dev/fixes-assembly

Small fixes/improvements to Assembly

Author: gonzalocasas

CHANGELOG.md

@@ -17,12 +17,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Added `Mesh.to_lines` method and tests.
 * Added `Data.guid` to JSON serialization.
 * Added `Data.guid` to pickle state.
+* Added `Assembly.find_by_key` to locate parts by key.
 
 ### Changed
 
 * Set `jinja >= 3.0` to dev dependencies to fix docs build error.
 * Fixed removing of collections for `compas_plotters`.
 * Fixed bug in `compas.robots.Configuration`.
+* Rebuild part index after deserialization in `Assembly`.
 
 ### Removed
 

src/compas/datastructures/__init__.py

@@ -23,6 +23,17 @@
     Part
 
 
+Exceptions
+==========
+
+.. autosummary::
+    :toctree: generated/
+    :nosignatures:
+
+    AssemblyError
+    FeatureError
+
+
 Functions
 =========
 
@@ -296,7 +307,9 @@
 
 from .assembly import (
     Assembly,
-    Part
+    Part,
+    AssemblyError,
+    FeatureError,
 )
 
 if not compas.IPY:
@@ -439,6 +452,8 @@
     # Assemblies
     'Assembly',
     'Part',
+    'AssemblyError',
+    'FeatureError',
 ]
 
 if not compas.IPY:

src/compas/datastructures/assembly/__init__.py

@@ -1,6 +1,8 @@
-from __future__ import print_function
 from __future__ import absolute_import
 from __future__ import division
+from __future__ import print_function
 
-from .assembly import Assembly  # noqa: F401 F403
-from .part import Part  # noqa: F401 F403
+from .exceptions import AssemblyError  # noqa: F401
+from .exceptions import FeatureError  # noqa: F401
+from .assembly import Assembly  # noqa: F401
+from .part import Part  # noqa: F401

src/compas/datastructures/assembly/assembly.py

@@ -30,7 +30,7 @@ class Assembly(Datastructure):
 
     def __init__(self, name=None, **kwargs):
         super(Assembly, self).__init__()
-        self.attributes = {'name': name or 'Assembly'}
+        self.attributes = {"name": name or "Assembly"}
         self.attributes.update(kwargs)
         self.graph = Graph()
         self._parts = {}
@@ -42,39 +42,43 @@ def __init__(self, name=None, **kwargs):
     @property
     def DATASCHEMA(self):
         import schema
-        return schema.Schema({
-            "attributes": dict,
-            "graph": Graph,
-        })
+
+        return schema.Schema(
+            {
+                "attributes": dict,
+                "graph": Graph,
+            }
+        )
 
     @property
     def JSONSCHEMANAME(self):
-        return 'assembly'
+        return "assembly"
 
     @property
     def data(self):
         data = {
-            'attributes': self.attributes,
-            'graph': self.graph.data,
+            "attributes": self.attributes,
+            "graph": self.graph.data,
         }
         return data
 
     @data.setter
     def data(self, data):
-        self.attributes.update(data['attributes'] or {})
-        self.graph.data = data['graph']
+        self.attributes.update(data["attributes"] or {})
+        self.graph.data = data["graph"]
+        self._parts = {part.guid: part.key for part in self.parts()}
 
     # ==========================================================================
     # properties
     # ==========================================================================
 
     @property
     def name(self):
-        return self.attributes.get('name') or self.__class__.__name__
+        return self.attributes.get("name") or self.__class__.__name__
 
     @name.setter
     def name(self, value):
-        self.attributes['name'] = value
+        self.attributes["name"] = value
 
     # ==========================================================================
     # customization
@@ -114,10 +118,10 @@ def add_part(self, part, key=None, **kwargs):
 
         """
         if part.guid in self._parts:
-            raise AssemblyError('Part already added to the assembly')
+            raise AssemblyError("Part already added to the assembly")
         key = self.graph.add_node(key=key, part=part, **kwargs)
         part.key = key
-        self._parts[part.guid] = part
+        self._parts[part.guid] = part.key
         return key
 
     def add_connection(self, a, b, **kwargs):
@@ -143,10 +147,11 @@ def add_connection(self, a, b, **kwargs):
             If `a` and/or `b` are not in the assembly.
 
         """
+        error_msg = "Both parts have to be added to the assembly before a connection can be created."
         if a.key is None or b.key is None:
-            raise AssemblyError('Both parts have to be added to the assembly before a connection can be created.')
+            raise AssemblyError(error_msg)
         if not self.graph.has_node(a.key) or not self.graph.has_node(b.key):
-            raise AssemblyError('Both parts have to be added to the assembly before a connection can be created.')
+            raise AssemblyError(error_msg)
         return self.graph.add_edge(a.key, b.key, **kwargs)
 
     def parts(self):
@@ -159,7 +164,7 @@ def parts(self):
 
         """
         for node in self.graph.nodes():
-            yield self.graph.node_attribute(node, 'part')
+            yield self.graph.node_attribute(node, "part")
 
     def connections(self, data=False):
         """Iterate over the connections between the parts.
@@ -194,4 +199,29 @@ def find(self, guid):
             or None if the part can't be found.
 
         """
-        return self._parts.get(guid)
+        key = self._parts.get(guid)
+
+        if key is None:
+            return None
+
+        return self.graph.node_attribute(key, "part")
+
+    def find_by_key(self, key):
+        """Find a part in the assembly by its key.
+
+        Parameters
+        ----------
+        key : int | str, optional
+            The identifier of the part in the assembly.
+
+        Returns
+        -------
+        :class:`~compas.datastructures.Part` | None
+            The identified part,
+            or None if the part can't be found.
+
+        """
+        if key not in self.graph.node:
+            return None
+
+        return self.graph.node_attribute(key, "part")

src/compas/datastructures/assembly/exceptions.py

@@ -1,11 +1,6 @@
-
 class AssemblyError(Exception):
     pass
 
 
-class FrameError(Exception):
-    pass
-
-
 class FeatureError(Exception):
     pass

tests/compas/datastructures/test_assembly.py

@@ -1,14 +1,17 @@
-from compas.datastructures import Assembly, assembly
+import pytest
+
+from compas.datastructures import Assembly
+from compas.datastructures import AssemblyError
 from compas.datastructures import Part
 
 
 def test_init():
-    assembly = Assembly(name='abc')
-    assert assembly.name == 'abc'
+    assembly = Assembly(name="abc")
+    assert assembly.name == "abc"
 
-    assembly = Assembly(attr1='value', attr2=3.14)
-    assert assembly.attributes['attr1'] == 'value'
-    assert assembly.attributes['attr2'] == 3.14
+    assembly = Assembly(attr1="value", attr2=3.14)
+    assert assembly.attributes["attr1"] == "value"
+    assert assembly.attributes["attr2"] == 3.14
 
 
 def test_add_parts():
@@ -20,6 +23,16 @@ def test_add_parts():
     assert len(list(assembly.parts())) == 3
 
 
+def test_add_duplicate_parts():
+    assembly = Assembly()
+
+    part = Part()
+    assembly.add_part(part)
+
+    with pytest.raises(AssemblyError):
+        assembly.add_part(part)
+
+
 def test_add_connections():
     assembly = Assembly()
     parts = [Part() for i in range(3)]
@@ -41,3 +54,24 @@ def test_find():
 
     assembly.add_part(part)
     assert assembly.find(part.guid) == part
+
+
+def test_find_by_key():
+    assembly = Assembly()
+    part = Part()
+    assembly.add_part(part, key=2)
+    assert assembly.find_by_key(2) == part
+
+    part = Part()
+    assembly.add_part(part, key="6")
+    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