From 38835a4092700ec54b584a1171115c2b626e0afa Mon Sep 17 00:00:00 2001
From: Dion Moult <dion@thinkmoult.com>
Date: Sat, 25 Jan 2025 14:11:42 +1100
Subject: [PATCH] See #6030. You can now select elements in groups recursively.

---
 .../bonsai/bim/module/group/__init__.py       |  3 +-
 .../bonsai/bim/module/group/operator.py       | 37 ++++++-------------
 src/bonsai/bonsai/bim/module/group/ui.py      |  2 +-
 .../ifcopenshell/util/element.py              |  6 ++-
 4 files changed, 18 insertions(+), 30 deletions(-)

diff --git a/src/bonsai/bonsai/bim/module/group/__init__.py b/src/bonsai/bonsai/bim/module/group/__init__.py
index 970c6cfbb8..8e69ee49d3 100644
--- a/src/bonsai/bonsai/bim/module/group/__init__.py
+++ b/src/bonsai/bonsai/bim/module/group/__init__.py
@@ -28,10 +28,9 @@ classes = (
     operator.EnableEditingGroup,
     operator.LoadGroups,
     operator.RemoveGroup,
-    operator.SelectGroupProducts,
+    operator.SelectGroupElements,
     operator.ToggleGroup,
     operator.UnassignGroup,
-    operator.SelectGroupElements,
     prop.ExpandedGroups,
     prop.Group,
     prop.BIMGroupProperties,
diff --git a/src/bonsai/bonsai/bim/module/group/operator.py b/src/bonsai/bonsai/bim/module/group/operator.py
index cf1d761bf5..4abf12f2ca 100644
--- a/src/bonsai/bonsai/bim/module/group/operator.py
+++ b/src/bonsai/bonsai/bim/module/group/operator.py
@@ -201,41 +201,26 @@ class UnassignGroup(bpy.types.Operator, tool.Ifc.Operator):
         ifcopenshell.api.group.unassign_group(tool.Ifc.get(), products=products, group=tool.Ifc.get().by_id(self.group))
 
 
-class SelectGroupProducts(bpy.types.Operator, tool.Ifc.Operator):
-    bl_idname = "bim.select_group_products"
-    bl_label = "Select Group Products"
-    bl_options = {"REGISTER", "UNDO"}
-    bl_description = "Select objects assigned to the selected group"
-    group: bpy.props.IntProperty()
-
-    def _execute(self, context):
-        self.file = IfcStore.get_file()
-        for obj in context.visible_objects:
-            obj.select_set(False)
-            element = tool.Ifc.get_entity(obj)
-            if not element:
-                continue
-            product_groups = [
-                r.RelatingGroup.id()
-                for r in getattr(element, "HasAssignments", []) or []
-                if r.is_a("IfcRelAssignsToGroup")
-            ]
-            if self.group in product_groups:
-                obj.select_set(True)
-        return {"FINISHED"}
-
-
 class SelectGroupElements(bpy.types.Operator):
     bl_idname = "bim.select_group_elements"
     bl_label = "Select Group elements"
     bl_options = {"REGISTER", "UNDO"}
+    bl_description = (
+        "Select objects assigned to the selected group and all nested groups\nALT + CLICK to exclude children"
+    )
     group: bpy.props.IntProperty()
+    is_recursive: bpy.props.BoolProperty(name="Is Recursive", default=True, options={"SKIP_SAVE"})
 
     @classmethod
     def poll(cls, context):
         return bool(tool.Ifc.get() and context.active_object)
 
+    def invoke(self, context, event):
+        self.is_recursive = not event.alt
+        return self.execute(context)
+
     def execute(self, context):
-        elements = tool.Drawing.get_group_elements(tool.Ifc.get().by_id(self.group))
-        tool.Spatial.select_products(elements)
+        tool.Spatial.select_products(
+            ifcopenshell.util.element.get_grouped_by(tool.Ifc.get().by_id(self.group), is_recursive=self.is_recursive)
+        )
         return {"FINISHED"}
diff --git a/src/bonsai/bonsai/bim/module/group/ui.py b/src/bonsai/bonsai/bim/module/group/ui.py
index f85b235258..dbd266a18f 100644
--- a/src/bonsai/bonsai/bim/module/group/ui.py
+++ b/src/bonsai/bonsai/bim/module/group/ui.py
@@ -55,7 +55,7 @@ class BIM_PT_groups(Panel):
         if (group := self.props.active_group) and (group_id := group.ifc_definition_id):
             row = self.layout.row(align=True)
             row.alignment = "RIGHT"
-            row.operator("bim.select_group_products", text="", icon="RESTRICT_SELECT_OFF").group = group_id
+            row.operator("bim.select_group_elements", text="", icon="RESTRICT_SELECT_OFF").group = group_id
             row.operator("bim.assign_group", text="", icon="FOLDER_REDIRECT").group = group_id
             row.operator("bim.enable_editing_group", text="", icon="GREASEPENCIL").group = group_id
             row.operator("bim.add_group", text="", icon="ADD").group = group_id
diff --git a/src/ifcopenshell-python/ifcopenshell/util/element.py b/src/ifcopenshell-python/ifcopenshell/util/element.py
index 168c1b19e0..bc2e28bc55 100644
--- a/src/ifcopenshell-python/ifcopenshell/util/element.py
+++ b/src/ifcopenshell-python/ifcopenshell/util/element.py
@@ -1000,7 +1000,9 @@ def get_decomposition(element: ifcopenshell.entity_instance, is_recursive=True)
     return results
 
 
-def get_grouped_by(element: ifcopenshell.entity_instance) -> list[ifcopenshell.entity_instance]:
+def get_grouped_by(
+    element: ifcopenshell.entity_instance, is_recursive: bool = True
+) -> list[ifcopenshell.entity_instance]:
     """Retrieves all subelements of an element based on the group.
 
     :param element: IfcGroup entity
@@ -1021,6 +1023,8 @@ def get_grouped_by(element: ifcopenshell.entity_instance) -> list[ifcopenshell.e
             related_objects = rel.RelatedObjects
             queue.extend(related_objects)
             results.extend(related_objects)
+        if not is_recursive:
+            break
     return results