From afbdedaee86f7b15187ded08619ffe77a5b27adb Mon Sep 17 00:00:00 2001
From: Dion Moult <dion@thinkmoult.com>
Date: Thu, 7 Jan 2021 21:02:46 +1100
Subject: [PATCH] WIP minor fixes to aggregate and spatial collection handling.
 See #1222.

---
 .../blenderbim/bim/module/aggregate/operator.py       |  9 +++------
 src/ifcblenderexport/blenderbim/bim/module/pset/ui.py | 11 ++++++++++-
 .../blenderbim/bim/module/root/operator.py            |  9 +++++++--
 src/ifcblenderexport/blenderbim/bim/prop.py           |  3 +++
 4 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/src/ifcblenderexport/blenderbim/bim/module/aggregate/operator.py b/src/ifcblenderexport/blenderbim/bim/module/aggregate/operator.py
index f5cf7659d7..7644e12699 100644
--- a/src/ifcblenderexport/blenderbim/bim/module/aggregate/operator.py
+++ b/src/ifcblenderexport/blenderbim/bim/module/aggregate/operator.py
@@ -1,5 +1,4 @@
 import bpy
-import ifcopenshell.util.schema
 import blenderbim.bim.module.aggregate.assign_object as assign_object
 from blenderbim.bim.ifc import IfcStore
 from blenderbim.bim.module.aggregate.data import Data
@@ -29,11 +28,9 @@ class AssignObject(bpy.types.Operator):
         Data.load(props.ifc_definition_id)
         bpy.ops.bim.disable_editing_aggregate(obj=related_object.name)
 
-        declaration = IfcStore.get_schema().declaration_by_name(product.is_a())
-        # TODO: we may not need this conditional if aggregates stop using collection instances
-        if ifcopenshell.util.schema.is_a(declaration, "IfcSpatialElement"):
-            related_collection = related_object.users_collection[0]
-            relating_collection = relating_object.users_collection[0]
+        related_collection = bpy.data.collections.get(related_object.name)
+        if related_collection:
+            relating_collection = bpy.data.collections.get(relating_object.name)
             self.remove_collection(bpy.context.scene.collection, related_collection)
             for collection in bpy.data.collections:
                 if collection == relating_collection:
diff --git a/src/ifcblenderexport/blenderbim/bim/module/pset/ui.py b/src/ifcblenderexport/blenderbim/bim/module/pset/ui.py
index 6aac43b0b0..43411fa14b 100644
--- a/src/ifcblenderexport/blenderbim/bim/module/pset/ui.py
+++ b/src/ifcblenderexport/blenderbim/bim/module/pset/ui.py
@@ -11,7 +11,16 @@ class BIM_PT_object_psets(Panel):
 
     @classmethod
     def poll(cls, context):
-        return context.active_object is not None and hasattr(context.active_object, "BIMObjectProperties")
+        if not context.active_object:
+            return False
+        props = context.active_object.BIMObjectProperties
+        if not props.ifc_definition_id:
+            return False
+        if props.ifc_definition_id not in Data.products:
+            Data.load(props.ifc_definition_id)
+        if not Data.products[props.ifc_definition_id]:
+            return False
+        return True
 
     def draw(self, context):
         props = context.active_object.BIMObjectProperties
diff --git a/src/ifcblenderexport/blenderbim/bim/module/root/operator.py b/src/ifcblenderexport/blenderbim/bim/module/root/operator.py
index aebd873f1f..8f15853a81 100644
--- a/src/ifcblenderexport/blenderbim/bim/module/root/operator.py
+++ b/src/ifcblenderexport/blenderbim/bim/module/root/operator.py
@@ -110,9 +110,14 @@ class AssignClass(bpy.types.Operator):
 
         bpy.ops.bim.add_representation(obj=obj.name)
 
-        if ifcopenshell.util.schema.is_a(self.declaration, "IfcElementType"):
+        if product.is_a("IfcElementType"):
             self.place_in_types_collection(obj)
-        elif ifcopenshell.util.schema.is_a(self.declaration, "IfcSpatialElement") or self.ifc_class == "IfcProject":
+        elif (
+            product.is_a("IfcSpatialElement")
+            or product.is_a("IfcSpatialStructureElement")
+            or product.is_a("IfcProject")
+            or product.is_a("IfcContext")
+        ):
             self.place_in_spatial_collection(obj)
         else:
             self.assign_potential_spatial_container(obj)
diff --git a/src/ifcblenderexport/blenderbim/bim/prop.py b/src/ifcblenderexport/blenderbim/bim/prop.py
index ebe645b45e..332818cacc 100644
--- a/src/ifcblenderexport/blenderbim/bim/prop.py
+++ b/src/ifcblenderexport/blenderbim/bim/prop.py
@@ -248,6 +248,7 @@ def refreshActiveDrawingIndex(self, context):
 
 def getIfcProducts(self, context):
     global products_enum
+    file = IfcStore.get_file()
     if len(products_enum) < 1:
         products_enum.extend(
             [
@@ -263,6 +264,8 @@ def getIfcProducts(self, context):
                 ]
             ]
         )
+        if file.schema == "IFC2X3":
+            products_enum[2] = ("IfcSpatialStructureElement", "IfcSpatialStructureElement", "")
     return products_enum