API now supports adding classifications by identifier instead of from a library

src/ifcopenshell-python/ifcopenshell/api/classification/add_reference.py

@@ -26,6 +26,8 @@ class Usecase:
         self.settings = {
             "product": None,
             "reference": None,
+            "identification": None,
+            "name": None,
             "classification": None,
             "is_lightweight": True,
         }
@@ -33,29 +35,34 @@ class Usecase:
             self.settings[key] = value
 
     def execute(self):
-        relating_classification = None
+        if self.settings["reference"]:
+            return self.add_from_library()
+        return self.add_from_identification()
 
+    def add_from_identification(self):
+        reference = self.get_existing_reference(self.settings["identification"])
+        if reference:
+            self.add_to_existing_relationship()
+        else:
+            reference = self.file.createIfcClassificationReference(
+                Identification=self.settings["identification"],
+                Name=self.settings["name"],
+                ReferencedSource=self.settings["classification"],
+            )
+            self.add_new_relationship(reference)
+        return reference
+
+    def add_from_library(self):
         if hasattr(self.settings["reference"], "ItemReference"):
             identification = self.settings["reference"].ItemReference  # IFC2X3
         else:
             identification = self.settings["reference"].Identification
 
-        for reference in self.file.by_type("IfcClassificationReference"):
-            if self.file.schema == "IFC2X3":
-                if reference.ItemReference == identification:
-                    relating_classification = reference
-                    break
-            else:
-                if reference.Identification == identification:
-                    relating_classification = reference
-                    break
+        reference = self.get_existing_reference(identification)
 
-        if relating_classification:
-            association = self.get_association(relating_classification)
-            related_objects = set(association.RelatedObjects)
-            related_objects.add(self.settings["product"])
-            association.RelatedObjects = list(related_objects)
-            return
+        if reference:
+            self.add_to_existing_relationship()
+            return reference
 
         migrator = ifcopenshell.util.schema.Migrator()
 
@@ -67,31 +74,46 @@ class Usecase:
                 c for c in self.file.by_type("IfcClassification") if c.Name == self.settings["classification"].Name
             ]
 
-        relating_classification = migrator.migrate(self.settings["reference"], self.file)
+        reference = migrator.migrate(self.settings["reference"], self.file)
 
         if self.settings["is_lightweight"]:
-            relating_classification.ReferencedSource = self.settings["classification"]
+            reference.ReferencedSource = self.settings["classification"]
             self.settings["reference"].ReferencedSource = old_referenced_source
         elif existing_classification:
             to_delete = set()
-            for traversed_reference in self.file.traverse(relating_classification):
+            for traversed_reference in self.file.traverse(reference):
                 if traversed_reference.ReferencedSource.is_a("IfcClassification"):
                     to_delete.add(traversed_reference.ReferencedSource)
                     traversed_reference.ReferencedSource = existing_classification[0]
                     break
             for element in to_delete:
                 self.file.remove(element)
+        self.add_new_relationship(reference)
 
+    def get_existing_reference(self, identification):
+        for reference in self.file.by_type("IfcClassificationReference"):
+            if self.file.schema == "IFC2X3":
+                if reference.ItemReference == identification:
+                    return reference
+            else:
+                if reference.Identification == identification:
+                    return reference
+
+    def add_new_relationship(self, reference):
         self.file.create_entity(
             "IfcRelAssociatesClassification",
-            **{
-                "GlobalId": ifcopenshell.guid.new(),
-                "RelatedObjects": [self.settings["product"]],
-                "RelatingClassification": relating_classification,
-            }
+            GlobalId=ifcopenshell.guid.new(),
+            RelatedObjects=[self.settings["product"]],
+            RelatingClassification=reference,
         )
 
-    def get_association(self, reference):
+    def add_to_existing_relationship(self, reference):
+        rel = self.get_rel_associates_classification(reference)
+        related_objects = set(rel.RelatedObjects)
+        related_objects.add(self.settings["product"])
+        rel.RelatedObjects = list(related_objects)
+
+    def get_rel_associates_classification(self, reference):
         if self.file.schema == "IFC2X3":
             for association in self.file.by_type("IfcRelAssociatesClassification"):
                 if association.RelatingClassification == reference: