Add status assignment to Statuses UI

Example - https://files.catbox.moe/c6ydj7.mp4

Still need to add tests and support for userdefined statuses.

src/bonsai/bonsai/bim/module/sequence/__init__.py

@@ -38,6 +38,7 @@ classes = (
     operator.AssignProduct,
     operator.AssignRecurrencePattern,
     operator.AssignSuccessor,
+    operator.AssignStatus,
     operator.AssignWorkSchedule,
     operator.Bonsai_DatePicker,
     operator.CalculateTaskDuration,

src/bonsai/bonsai/bim/module/sequence/data.py

@@ -32,6 +32,7 @@ def refresh():
     TaskICOMData.is_loaded = False
     WorkScheduleData.is_loaded = False
     AnimationColorSchemeData.is_loaded = False
+    StatusData.is_loaded = False
 
 
 class SequenceData:
@@ -421,3 +422,25 @@ class AnimationColorSchemeData:
             except:
                 pass
         return [(str(g.id()), g.Name or "Unnamed", "") for g in sorted(results, key=lambda x: x.Name or "Unnamed")]
+
+
+class StatusData:
+    data: dict[str, Any] = {}
+    is_loaded = False
+
+    @classmethod
+    def load(cls) -> None:
+        cls.is_loaded = True
+        cls.data = {
+            "statuses_with_elements": cls.statuses_with_elements(),
+        }
+
+    @classmethod
+    def statuses_with_elements(cls) -> set[str]:
+        statuses = ["No Status"]
+        statuses.extend(tool.Sequence.ELEMENT_STATUSES)
+        statuses_used: set[str] = set()
+        for status in statuses:
+            if tool.Sequence.get_elements_by_status(status):
+                statuses_used.add(status)
+        return statuses_used

src/bonsai/bonsai/bim/module/sequence/operator.py

@@ -18,22 +18,26 @@
 
 # pyright: reportUnnecessaryTypeIgnoreComment=error
 
-import os
+import types
+
+from collections import Counter
+from functools import cache
 
 import bpy
-import json
 import time
 import calendar
-import isodate
+import ifcopenshell.api.pset
+import ifcopenshell.util.element
+import bonsai.bim.schema
 import bonsai.core.sequence as core
 import bonsai.tool as tool
-import bonsai.bim.module.sequence.helper as helper
-import ifcopenshell.util.sequence
-import ifcopenshell.util.selector
 from datetime import datetime
 from dateutil import parser, relativedelta
 from bpy_extras.io_utils import ImportHelper, ExportHelper
-from typing import get_args, TYPE_CHECKING, assert_never
+from typing import Union, get_args, TYPE_CHECKING, assert_never
+
+if TYPE_CHECKING:
+    from bpy.stub_internal.rna_enums import OperatorReturnItems
 
 
 class EnableStatusFilters(bpy.types.Operator):
@@ -48,25 +52,33 @@ class EnableStatusFilters(bpy.types.Operator):
 
         props.statuses.clear()
 
-        statuses = set()
+        statuses_used: Counter[str] = Counter()
+        user_defined_statuses: set[str] = set()
         for element in tool.Ifc.get().by_type("IfcPropertyEnumeratedValue"):
             if element.Name == "Status":
-                if element.PartOfPset and isinstance(element.EnumerationValues, tuple):
+                enum_values = element.EnumerationValues
+                if element.PartOfPset and isinstance(enum_values, tuple):
                     pset = element.PartOfPset[0]
-                    if pset.Name.startswith("Pset_") and pset.Name.endswith("Common"):
-                        statuses.update(element.EnumerationValues)
-                    elif pset.Name == "EPset_Status":  # Our secret sauce
-                        statuses.update(element.EnumerationValues)
+                    pset_name: str = pset.Name
+                    if pset_name.startswith("Pset_") and pset_name.endswith("Common"):
+                        statuses_used.update([s.wrappedValue for s in enum_values])
+                    elif pset_name == "EPset_Status":  # Our secret sauce
+                        statuses_used.update([s.wrappedValue for s in enum_values])
             elif element.Name == "UserDefinedStatus":
-                statuses.add(element.NominalValue)
+                status: str = element.NominalValue.wrappedValue
+                statuses_used[element.NominalValue.wrappedValue] += 1
+                user_defined_statuses.add(status)
 
-        statuses = ["No Status"] + sorted([s.wrappedValue for s in statuses])
+        statuses = ["No Status"]
+        statuses.extend(tool.Sequence.ELEMENT_STATUSES)
+        statuses.extend(user_defined_statuses)
 
         for status in statuses:
             new = props.statuses.add()
             new.name = status
             if new.name in hidden_statuses:
                 new.is_visible = False
+            new.has_elements = bool(statuses_used[status])
 
         visible_statuses = {s.name for s in props.statuses if s.is_visible}
         tool.Sequence.set_visibility_by_status(visible_statuses)
@@ -124,19 +136,140 @@ class SelectStatusFilter(bpy.types.Operator):
     bl_label = "Select Status Filter"
     bl_description = "Select elements with currently selected status"
     bl_options = {"REGISTER", "UNDO"}
-    name: bpy.props.StringProperty()
 
-    def execute(self, context):
-        query = f"IfcProduct, /Pset_.*Common/.Status={self.name} + IfcProduct, EPset_Status.Status={self.name}"
-        if self.name == "No Status":
-            query = f"IfcProduct, /Pset_.*Common/.Status=NULL, EPset_Status.Status=NULL"
-        for element in ifcopenshell.util.selector.filter_elements(tool.Ifc.get(), query):
+    status: bpy.props.StringProperty()  # pyright: ignore[reportRedeclaration]
+
+    if TYPE_CHECKING:
+        status: tool.Sequence.ElementStatusUI
+
+    def execute(self, context) -> set["OperatorReturnItems"]:
+        for element in tool.Sequence.get_elements_by_status(self.status):
             obj = tool.Ifc.get_object(element)
-            if obj:
+            if isinstance(obj, bpy.types.Object):
                 obj.select_set(True)
         return {"FINISHED"}
 
 
+class AssignStatus(bpy.types.Operator, tool.Ifc.Operator):
+    bl_idname = "bim.assign_status"
+    bl_label = "Assign Status"
+    bl_description = "Assign status to the selected elements.\n\nAlt+CLICK to unassign the status."
+    bl_options = {"REGISTER", "UNDO"}
+
+    should_override_previous_status: bpy.props.BoolProperty(  # pyright: ignore[reportRedeclaration]
+        name="Override Previous Status",
+        description=(
+            "Whether assigning new status should override previous one.\n\n"
+            "IFC allows storing multiple statuses for the same element. "
+            "This option can be disabled to take advantage of that."
+        ),
+        default=True,
+    )
+    status: bpy.props.StringProperty()  # pyright: ignore[reportRedeclaration]
+    should_unassign_status: bpy.props.BoolProperty(  # pyright: ignore[reportRedeclaration]
+        options={"SKIP_SAVE"},
+    )
+
+    if TYPE_CHECKING:
+        should_override_previous_status: bool
+        status: str
+        should_unassign_status: bool
+
+    def invoke(self, context, event):
+        self.should_unassign_status = event.alt
+        return self.execute(context)
+
+    def _execute(self, context):
+        # TODO: UserDefinedStatus
+        if self.status not in tool.Sequence.ELEMENT_STATUSES:
+            self.report({"ERROR"}, "Assigning user defined statuses or 'No Status' is not yet supported.")
+            return {"CANCELLED"}
+
+        EPSET_NAME = "EPset_Status"
+        elements_changed = 0
+        ifc_file = tool.Ifc.get()
+
+        @cache
+        def get_common_pset_name(element: ifcopenshell.entity_instance) -> Union[str, None]:
+            templates = bonsai.bim.schema.ifc.psetqto.get_applicable(
+                element.is_a(),
+                pset_only=True,
+                schema=tool.Ifc.get_schema(),
+            )
+            for template in templates:
+                template_name = template.Name
+                if template_name.startswith("Pset_") and template_name.endswith("Common"):
+                    return template_name
+
+        for obj in tool.Blender.get_selected_objects():
+            if not (element := tool.Ifc.get_entity(obj)) or not element.is_a("IfcProduct"):
+                continue
+
+            psets = ifcopenshell.util.element.get_psets(element, psets_only=True)
+            common_pset_name = get_common_pset_name(element)
+
+            existing_psets = [
+                pset_name for pset_name in psets if pset_name == EPSET_NAME or pset_name == common_pset_name
+            ]
+            assert len(existing_psets) < 3
+            # Common pset comes first.
+            existing_psets.sort(key=lambda x: x == EPSET_NAME)
+
+            if not existing_psets:
+                if self.should_unassign_status:
+                    continue
+                pset_name = common_pset_name or EPSET_NAME
+                pset = ifcopenshell.api.pset.add_pset(ifc_file, element, pset_name)
+                ifcopenshell.api.pset.edit_pset(ifc_file, pset, properties={"Status": [self.status]})
+                elements_changed += 1
+                continue
+
+            pset_changed = False
+            for pset_i, pset_name in enumerate(existing_psets):
+                pset_data = psets[pset_name]
+                # None is kind of theoretical.
+                status_data: Union[list[str], None, types.EllipsisType]
+                status_data = pset_data.get("Status", ...)
+
+                if self.should_unassign_status:
+                    if status_data is ... or not status_data:
+                        continue
+                    # Already unassigned.
+                    if self.status not in status_data:
+                        continue
+                    status_data.remove(self.status)
+
+                else:
+                    if status_data is None or status_data is ...:
+                        status_data = [self.status]
+                    elif self.status in status_data:
+                        # Already assigned.
+                        continue
+                    else:
+                        if self.should_override_previous_status:
+                            status_data = [self.status]
+                        else:
+                            status_data.append(self.status)
+
+                if pset_i > 0:
+                    # Try to maintain status in just 1 pset.
+                    if not self.should_unassign_status:
+                        status_data.remove(self.status)
+
+                ifcopenshell.api.pset.edit_pset(
+                    ifc_file, pset=ifc_file.by_id(pset_data["id"]), properties={"Status": status_data}
+                )
+                pset_changed = True
+            elements_changed += pset_changed
+
+        self.report(
+            {"INFO"},
+            f"Status '{self.status}' "
+            f"{'un' * self.should_unassign_status}assigned {'from' if self.should_unassign_status else 'to'} "
+            f"{elements_changed} elements.",
+        )
+
+
 class AddWorkPlan(bpy.types.Operator, tool.Ifc.Operator):
     bl_idname = "bim.add_work_plan"
     bl_label = "Add Work Plan"

src/bonsai/bonsai/bim/module/sequence/prop.py

@@ -41,7 +41,7 @@ from bpy.props import (
     FloatVectorProperty,
     CollectionProperty,
 )
-from typing import TYPE_CHECKING, Literal, get_args
+from typing import TYPE_CHECKING, Literal, Union, get_args
 
 
 def getTaskColumns(self, context):
@@ -429,13 +429,13 @@ class BIMWorkPlanProperties(PropertyGroup):
 
 
 class IFCStatus(PropertyGroup):
-    name: StringProperty(name="Name")
-    is_visible: BoolProperty(
+    name: StringProperty()  # pyright: ignore[reportRedeclaration]
+    is_visible: BoolProperty(  # pyright: ignore[reportRedeclaration]
         name="Is Visible", default=True, update=lambda x, y: (None, bpy.ops.bim.activate_status_filters())[0]
     )
 
     if TYPE_CHECKING:
-        name: str
+        name: tool.Sequence.ElementStatusUI
         is_visible: bool
 
 

src/bonsai/bonsai/bim/module/sequence/ui.py

@@ -29,6 +29,7 @@ from bonsai.bim.module.sequence.data import (
     SequenceData,
     TaskICOMData,
     AnimationColorSchemeData,
+    StatusData,
 )
 from typing import Any, Optional, TYPE_CHECKING
 
@@ -59,16 +60,25 @@ class BIM_PT_status(Panel):
             row.operator("bim.enable_status_filters", icon="GREASEPENCIL")
             return
 
+        if not StatusData.is_loaded:
+            StatusData.load()
+
         row = self.layout.row(align=True)
-        row.label(text="Statuses found in the project:")
+        row.label(text="Elements Statuses:")
         row.operator("bim.activate_status_filters", icon="FILE_REFRESH", text="")
         row.operator("bim.disable_status_filters", icon="CANCEL", text="")
 
+        box = self.layout.box()
         for status in self.props.statuses:
-            row = self.layout.row(align=True)
+            row = box.row(align=True)
             row.label(text=status.name)
+            if status.name in StatusData.data["statuses_with_elements"]:
+                row.label(text="", icon="ASSET_MANAGER")
             row.prop(status, "is_visible", text="", emboss=False, icon="HIDE_OFF" if status.is_visible else "HIDE_ON")
-            row.operator("bim.select_status_filter", icon="RESTRICT_SELECT_OFF", text="").name = status.name
+            row.operator("bim.select_status_filter", icon="RESTRICT_SELECT_OFF", text="").status = status.name
+            row.operator("bim.assign_status", icon="BRUSH_DATA", text="").status = status.name
+
+        # TODO: also add a prop to add custom userdefined status.
 
 
 class BIM_PT_work_plans(Panel):

src/bonsai/bonsai/tool/sequence.py

@@ -1840,15 +1840,29 @@ class Sequence(bonsai.core.tool.Sequence):
             return isodate.datetime_isoformat(datetime_)
         return isodate.date_isoformat(datetime_)
 
+    ElementStatus = Literal["NEW", "EXISTING", "DEMOLISH", "TEMPORARY", "OTHER", "NOTKNOWN", "UNSET"]
+    ElementStatusUI = Union[ElementStatus, Literal["No Status"], str]
+    """Also allows UserDefinedStatus from EPset."""
+
+    ELEMENT_STATUSES = ("NEW", "EXISTING", "DEMOLISH", "TEMPORARY", "OTHER", "NOTKNOWN", "UNSET")
+
+    @classmethod
+    def get_status_query(cls, status: ElementStatusUI) -> str:
+        if status == "No Status":
+            return f"IfcProduct, /Pset_.*Common/.Status=NULL, EPset_Status.Status=NULL"
+        return f"IfcProduct, /Pset_.*Common/.Status={status} + IfcProduct, EPset_Status.Status={status}"
+
+    @classmethod
+    def get_elements_by_status(cls, status: ElementStatusUI) -> set[ifcopenshell.entity_instance]:
+        query = cls.get_status_query(status)
+        return ifcopenshell.util.selector.filter_elements(tool.Ifc.get(), query)
+
     @classmethod
     def set_visibility_by_status(cls, visible_statuses: set[str]) -> None:
         assert bpy.context.view_layer
         query = []
         for name in visible_statuses:
-            if name == "No Status":
-                q = f"IfcProduct, /Pset_.*Common/.Status=NULL, EPset_Status.Status=NULL"
-            else:
-                q = f"IfcProduct, /Pset_.*Common/.Status={name} + IfcProduct, EPset_Status.Status={name}"
+            q = cls.get_status_query(name)
             query.append(q)
         query = " + ".join(query)
 

src/ifcopenshell-python/ifcopenshell/util/pset.py

@@ -78,6 +78,7 @@ class PsetQto:
         qto_only=False,
         schema: ifcopenshell.util.schema.IFC_SCHEMA = "IFC4",
     ) -> list[entity_instance]:
+        """Get applicable property set templates."""
         any_class = not ifc_class
         entity = None
         if not any_class: