⚫ black

Author: gonzalocasas

src/compas/datastructures/assembly/assembly.py

@@ -41,39 +41,42 @@ 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"]
 
     # ==========================================================================
     # 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
@@ -143,10 +146,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 +163,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.

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)]