From 4c87ea321dc2eba578619539ceb1f16ed68412a5 Mon Sep 17 00:00:00 2001
From: Dion Moult <dion@thinkmoult.com>
Date: Fri, 29 Jan 2021 15:15:14 +1100
Subject: [PATCH] Improve UX of assigning classes to multiple objects. See
 #1222.

---
 .../blenderbim/bim/module/model/__init__.py   |  2 +
 .../blenderbim/bim/module/model/pie.py        | 40 ++++++++++++-------
 .../blenderbim/bim/module/root/ui.py          |  1 -
 3 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/src/ifcblenderexport/blenderbim/bim/module/model/__init__.py b/src/ifcblenderexport/blenderbim/bim/module/model/__init__.py
index 8e097c91b3..23f5bbb353 100644
--- a/src/ifcblenderexport/blenderbim/bim/module/model/__init__.py
+++ b/src/ifcblenderexport/blenderbim/bim/module/model/__init__.py
@@ -17,6 +17,8 @@ classes = (
     pie.AssignIfcStair,
     pie.AssignIfcDoor,
     pie.AssignIfcWindow,
+    pie.AssignIfcColumn,
+    pie.AssignIfcBeam,
     pie.VIEW3D_MT_PIE_bim,
     pie.VIEW3D_MT_PIE_bim_class,
 )
diff --git a/src/ifcblenderexport/blenderbim/bim/module/model/pie.py b/src/ifcblenderexport/blenderbim/bim/module/model/pie.py
index 910e9d0b7b..25767436e0 100644
--- a/src/ifcblenderexport/blenderbim/bim/module/model/pie.py
+++ b/src/ifcblenderexport/blenderbim/bim/module/model/pie.py
@@ -15,9 +15,7 @@ class AssignIfcWall(bpy.types.Operator):
     bl_label = "IfcWall"
 
     def execute(self, context):
-        for obj in context.selected_objects:
-            if not obj.BIMObjectProperties.ifc_definition_id:
-                bpy.ops.bim.assign_class(obj=obj.name, ifc_class="IfcWall")
+        bpy.ops.bim.assign_class(ifc_class="IfcWall")
         return {"FINISHED"}
 
 
@@ -26,9 +24,7 @@ class AssignIfcSlab(bpy.types.Operator):
     bl_label = "IfcSlab"
 
     def execute(self, context):
-        for obj in context.selected_objects:
-            if not obj.BIMObjectProperties.ifc_definition_id:
-                bpy.ops.bim.assign_class(obj=obj.name, ifc_class="IfcSlab")
+        bpy.ops.bim.assign_class(ifc_class="IfcSlab")
         return {"FINISHED"}
 
 
@@ -37,9 +33,7 @@ class AssignIfcStair(bpy.types.Operator):
     bl_label = "IfcStair"
 
     def execute(self, context):
-        for obj in context.selected_objects:
-            if not obj.BIMObjectProperties.ifc_definition_id:
-                bpy.ops.bim.assign_class(obj=obj.name, ifc_class="IfcStair")
+        bpy.ops.bim.assign_class(ifc_class="IfcStair")
         return {"FINISHED"}
 
 
@@ -48,9 +42,7 @@ class AssignIfcDoor(bpy.types.Operator):
     bl_label = "IfcDoor"
 
     def execute(self, context):
-        for obj in context.selected_objects:
-            if not obj.BIMObjectProperties.ifc_definition_id:
-                bpy.ops.bim.assign_class(obj=obj.name, ifc_class="IfcDoor")
+        bpy.ops.bim.assign_class(ifc_class="IfcDoor")
         return {"FINISHED"}
 
 
@@ -59,9 +51,25 @@ class AssignIfcWindow(bpy.types.Operator):
     bl_label = "IfcWindow"
 
     def execute(self, context):
-        for obj in context.selected_objects:
-            if not obj.BIMObjectProperties.ifc_definition_id:
-                bpy.ops.bim.assign_class(obj=obj.name, ifc_class="IfcWindow")
+        bpy.ops.bim.assign_class(ifc_class="IfcWindow")
+        return {"FINISHED"}
+
+
+class AssignIfcColumn(bpy.types.Operator):
+    bl_idname = "bim.assign_ifc_column"
+    bl_label = "IfcColumn"
+
+    def execute(self, context):
+        bpy.ops.bim.assign_class(ifc_class="IfcColumn")
+        return {"FINISHED"}
+
+
+class AssignIfcBeam(bpy.types.Operator):
+    bl_idname = "bim.assign_ifc_beam"
+    bl_label = "IfcBeam"
+
+    def execute(self, context):
+        bpy.ops.bim.assign_class(ifc_class="IfcBeam")
         return {"FINISHED"}
 
 
@@ -123,3 +131,5 @@ class VIEW3D_MT_PIE_bim_class(bpy.types.Menu):
         pie.operator("bim.assign_ifc_stair")
         pie.operator("bim.assign_ifc_door")
         pie.operator("bim.assign_ifc_window")
+        pie.operator("bim.assign_ifc_column")
+        pie.operator("bim.assign_ifc_beam")
diff --git a/src/ifcblenderexport/blenderbim/bim/module/root/ui.py b/src/ifcblenderexport/blenderbim/bim/module/root/ui.py
index def15004cb..4a8345dde0 100644
--- a/src/ifcblenderexport/blenderbim/bim/module/root/ui.py
+++ b/src/ifcblenderexport/blenderbim/bim/module/root/ui.py
@@ -42,7 +42,6 @@ class BIM_PT_class(Panel):
             self.draw_class_dropdowns()
             row = self.layout.row(align=True)
             op = row.operator("bim.assign_class")
-            op.obj = context.active_object.name
             op.ifc_class = bpy.context.scene.BIMRootProperties.ifc_class
             op.predefined_type = bpy.context.scene.BIMRootProperties.ifc_predefined_type
             op.userdefined_type = bpy.context.scene.BIMRootProperties.ifc_userdefined_type