Refactor cost schedule CSV filepath to use document references

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

@@ -541,11 +541,10 @@ class SelectCostScheduleProducts(bpy.types.Operator):
         )
         return {"FINISHED"}
 
-
 class ImportCostScheduleCsv(bpy.types.Operator, ImportHelper, tool.Ifc.Operator):
     bl_idname = "bim.import_cost_schedule_csv"
     bl_label = "Import Cost Schedule CSV"
-    bl_description = "Import cost schdule from the provided .csv file."
+    bl_description = "Import cost schedule from the provided .csv file."
     bl_options = {"REGISTER", "UNDO"}
     filename_ext = ".csv"
     filter_glob: bpy.props.StringProperty(default="*.csv", options={"HIDDEN"})
@@ -568,11 +567,27 @@ class ImportCostScheduleCsv(bpy.types.Operator, ImportHelper, tool.Ifc.Operator)
 class RefreshCostScheduleCsv(bpy.types.Operator, tool.Ifc.Operator):
     bl_idname = "bim.refresh_cost_schedule_csv"
     bl_label = "Refresh Cost Schedule CSV"
+    bl_description = "Refresh cost schedule data from the associated CSV file"
     bl_options = {"REGISTER", "UNDO"}
 
     @classmethod
+    def poll(cls, context):
+        props = tool.Cost.get_cost_props()
+        if not props.active_cost_schedule_id:
+            cls.poll_message_set("No active cost schedule")
+            return False
+        
+        filepath = tool.Cost.get_cost_schedule_csv_filepath(props.active_cost_schedule_id)
+        if not filepath:
+            cls.poll_message_set("No CSV file associated with this cost schedule")
+            return False
+            
+        return True
+
     def _execute(self, context):
-        core.refresh_cost_schedule_csv(tool.Ifc, tool.Cost)
+        tool.Cost.delete_all_cost_items()
+        tool.Cost.refresh_cost_schedule_csv()
+        tool.Cost.load_cost_schedule_tree()
         return {"FINISHED"}
 
 

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

@@ -76,16 +76,12 @@ class BIM_PT_cost_schedules(Panel):
             col = row0.column()
             col.label(text="Linked CSV:")
             row_1 = col.row(align=True)
-            if self.props.active_cost_schedule_id in [item.cost_schedule_id for item in self.props.cost_schedule_files]:
-                file = next(
-                    (
-                        item.csv_filepath
-                        for item in self.props.cost_schedule_files
-                        if item.cost_schedule_id == self.props.active_cost_schedule_id
-                    ),
-                    None,
-                )
-                row_1.label(text=file)
+            
+            # Get filepath from document reference instead of props
+            file_path = tool.Cost.get_cost_schedule_csv_filepath(self.props.active_cost_schedule_id)
+            
+            if file_path:
+                row_1.label(text=file_path)
                 row_1.operator("bim.refresh_cost_schedule_csv", icon="FILE_REFRESH", text="")
             else:
                 row_1.label(text="No CSV file found")

src/bonsai/bonsai/core/cost.py

@@ -349,7 +349,6 @@ def import_cost_schedule_csv(
 def add_csv_filepath(cost: type[tool.Cost], file_path: str, is_schedule_of_rates: bool, cost_schedule) -> None:
     cost.add_csv_filepath(file_path, is_schedule_of_rates, cost_schedule)
 
-
 def remove_csv_filepath(cost: type[tool.Cost], cost_schedule) -> None:
     cost.remove_csv_filepath(cost_schedule)
 

src/bonsai/bonsai/tool/cost.py

@@ -590,39 +590,68 @@ class Cost(bonsai.core.tool.Cost):
         return csv2ifc.cost_schedule
 
     @classmethod
-    def add_csv_filepath(
-        cls,
-        file_path: Optional[str] = None,
-        is_schedule_of_rates: bool = False,
-        cost_schedule: ifcopenshell.entity_instance = None,
-    ) -> None:
-        if not file_path:
+    def add_csv_filepath(cls, file_path: str, is_schedule_of_rates: bool, cost_schedule) -> None:
+        """Store CSV filepath as a document reference in the IFC file."""
+        if not file_path or not cost_schedule:
             return
 
-        props = cls.get_cost_props()
-        if not props.active_cost_schedule_id in [item.cost_schedule_id for item in props.cost_schedule_files]:
-            item = props.cost_schedule_files.add()
-            item.cost_schedule_id = cost_schedule.id()
-            item.csv_filepath = file_path
+        ifc_file = tool.Ifc.get()
+        cost_docs_document = next(
+            (
+                document
+                for document in ifc_file.by_type("IfcDocumentInformation")
+                if document.Name == "BBIM_Cost_Documents"
+            ),
+            None,
+        )
+
+        if not cost_docs_document:
+            cost_docs_document = ifcopenshell.api.document.add_information(ifc_file)
+            cost_docs_document.Name = "BBIM_Cost_Documents"
+            cost_docs_document.Description = "Bonsai internal document containing references to cost CSV files"
+
+        # Create a reference with the filepath
+        reference = ifcopenshell.api.document.add_reference(ifc_file, cost_docs_document)
+        reference.Location = file_path
+        
+        # Store the cost schedule ID in the reference description
+        reference.Description = f"Cost Schedule ID: {cost_schedule.id()}"
+        
+        # If it's a schedule of rates, store that info too
+        if is_schedule_of_rates:
+            reference.Identification = "SCHEDULE_OF_RATES"
         else:
-            return
+            reference.Identification = "COST_SCHEDULE"
 
     @classmethod
     def remove_csv_filepath(cls, cost_schedule: ifcopenshell.entity_instance = None) -> None:
+        """Remove CSV filepath reference from IFC file."""
         if not cost_schedule:
             return
 
-        props = cls.get_cost_props()
-        cost_schedule_id = cost_schedule.id()
-        if cost_schedule_id in [item.cost_schedule_id for item in props.cost_schedule_files]:
-            for i, item in enumerate(props.cost_schedule_files):
-                if item.cost_schedule_id == cost_schedule_id:
-                    props.cost_schedule_files.remove(i)
-                    print(f"Cost schedule id={cost_schedule_id} csv filepath correctly removed")
-                    return
-        else:
+        ifc_file = tool.Ifc.get()
+        cost_docs_document = next(
+            (
+                document
+                for document in ifc_file.by_type("IfcDocumentInformation")
+                if document.Name == "BBIM_Cost_Documents"
+            ),
+            None,
+        )
+
+        if not cost_docs_document:
             return
 
+        cost_schedule_id = cost_schedule.id()
+        references = tool.Document.get_document_references(cost_docs_document)
+        
+        for reference in references:
+            # Check if this reference is for the given cost schedule
+            if reference.Description and f"Cost Schedule ID: {cost_schedule_id}" in reference.Description:
+                ifcopenshell.api.document.remove_reference(ifc_file, reference)
+                print(f"Cost schedule id={cost_schedule_id} csv filepath correctly removed")
+                return
+
     @classmethod
     def delete_all_cost_items(cls):
         cost_schedule = tool.Cost.get_active_cost_schedule()
@@ -632,25 +661,75 @@ class Cost(bonsai.core.tool.Cost):
             ifcopenshell.api.cost.remove_cost_item(tool.Ifc.get(), cost_item=item)
             tool.Cost.clean_up_cost_item_tree(cost_item_id)
 
+    @classmethod
+    def is_schedule_of_rates_csv(cls, cost_schedule_id: int) -> bool:
+        """Check if a cost schedule is a schedule of rates based on document references."""
+        ifc_file = tool.Ifc.get()
+        cost_docs_document = next(
+            (
+                document
+                for document in ifc_file.by_type("IfcDocumentInformation")
+                if document.Name == "BBIM_Cost_Documents"
+            ),
+            None,
+        )
+
+        if not cost_docs_document:
+            return False
+
+        references = tool.Document.get_document_references(cost_docs_document)
+        
+        for reference in references:
+            if reference.Description and f"Cost Schedule ID: {cost_schedule_id}" in reference.Description:
+                return reference.Identification == "SCHEDULE_OF_RATES"
+                
+        return False
+
+    @classmethod
+    def get_cost_schedule_csv_filepath(cls, cost_schedule_id: int) -> Optional[str]:
+        """Get CSV filepath for a cost schedule from document references."""
+        ifc_file = tool.Ifc.get()
+        cost_docs_document = next(
+            (
+                document
+                for document in ifc_file.by_type("IfcDocumentInformation")
+                if document.Name == "BBIM_Cost_Documents"
+            ),
+            None,
+        )
+
+        if not cost_docs_document:
+            return None
+
+        references = tool.Document.get_document_references(cost_docs_document)
+        
+        for reference in references:
+            if reference.Description and f"Cost Schedule ID: {cost_schedule_id}" in reference.Description:
+                return reference.Location
+                
+        return None
+
+
     @classmethod
     def refresh_cost_schedule_csv(cls):
+        """Refresh cost schedule from CSV file stored in document references."""
         from ifc5d.csv2ifc import Csv2Ifc
 
         props = cls.get_cost_props()
         cost_schedule_id = props.active_cost_schedule_id
-        file_path = next(
-            (item.csv_filepath for item in props.cost_schedule_files if item.cost_schedule_id == cost_schedule_id), None
-        )
+        file_path = cls.get_cost_schedule_csv_filepath(cost_schedule_id)
+        
         if not file_path:
             return
 
         cost_schedule = tool.Ifc.get_entity_by_id(cost_schedule_id)
+        is_schedule_of_rates = cls.is_schedule_of_rates_csv(cost_schedule_id)
 
         csv2ifc = Csv2Ifc()
         csv2ifc.csv = file_path
         csv2ifc.file = tool.Ifc.get()
         csv2ifc.cost_schedule = cost_schedule
-        csv2ifc.is_schedule_of_rates = False
+        csv2ifc.is_schedule_of_rates = is_schedule_of_rates
         csv2ifc.refresh()
 
         print("Csv file correctly refreshed")
@@ -1032,3 +1111,51 @@ class Cost(bonsai.core.tool.Cost):
         if results["quantity_type"] == "IfcQuantityCount":
             results["unit_symbol"] = "U"
         return results
+
+    @classmethod
+    def get_cost_schedule_documents(cls) -> Union[ifcopenshell.entity_instance, None]:
+        """Get the document information entity that stores CSV references."""
+        for document in tool.Ifc.get().by_type("IfcDocumentInformation"):
+            if document.Name == "BBIM_Cost_Documents":
+                return document
+        return None
+
+    @classmethod
+    def add_csv_filepath(
+        cls,
+        file_path: Optional[str] = None,
+        is_schedule_of_rates: bool = False,
+        cost_schedule: ifcopenshell.entity_instance = None,
+    ) -> None:
+        """Store CSV filepath as a document reference in the IFC file."""
+        if not file_path or not cost_schedule:
+            return
+
+        ifc_file = tool.Ifc.get()
+        cost_docs_document = next(
+            (
+                document
+                for document in ifc_file.by_type("IfcDocumentInformation")
+                if document.Name == "BBIM_Cost_Documents"
+            ),
+            None,
+        )
+
+        if not cost_docs_document:
+            cost_docs_document = ifcopenshell.api.document.add_information(ifc_file)
+            cost_docs_document.Name = "BBIM_Cost_Documents"
+            cost_docs_document.Description = "Bonsai internal document containing references to cost CSV files"
+
+        # Create a reference with the filepath
+        reference = ifcopenshell.api.document.add_reference(ifc_file, cost_docs_document)
+        reference.Location = file_path
+        
+        # Store the cost schedule ID in the reference description
+        reference.Description = f"Cost Schedule ID: {cost_schedule.id()}"
+        
+        # If it's a schedule of rates, store that info too
+        if is_schedule_of_rates:
+            reference.Identification = "SCHEDULE_OF_RATES"
+        else:
+            reference.Identification = "COST_SCHEDULE"
+