From deb6e3779e48d714c541f0e4fd2b09d390356b60 Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Thu, 17 Apr 2025 11:17:37 +0500
Subject: [PATCH] ifc5d - refactor

---
 src/ifc5d/ifc5d/csv2ifc.py          |  3 +++
 src/ifc5d/ifc5d/ifc5Dspreadsheet.py | 32 +++++++++++++++--------------
 2 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/src/ifc5d/ifc5d/csv2ifc.py b/src/ifc5d/ifc5d/csv2ifc.py
index f53f9f58bf..8b2c72fc5a 100644
--- a/src/ifc5d/ifc5d/csv2ifc.py
+++ b/src/ifc5d/ifc5d/csv2ifc.py
@@ -165,6 +165,7 @@ class Csv2Ifc:
     def create_ifc(self) -> None:
         if not self.file:
             self.create_boilerplate_ifc()
+        assert self.file
         if not self.cost_schedule:
             self.cost_schedule = ifcopenshell.api.cost.add_cost_schedule(self.file, name="CSV Import")
             if self.is_schedule_of_rates:
@@ -179,6 +180,7 @@ class Csv2Ifc:
             self.create_cost_item(cost_item, parent)
 
     def create_cost_item(self, cost_item: CostItem, parent: Optional[ifcopenshell.entity_instance] = None) -> None:
+        assert self.file
         if parent is None:
             cost_item["ifc"] = ifcopenshell.api.cost.add_cost_item(self.file, cost_schedule=self.cost_schedule)
         else:
@@ -265,6 +267,7 @@ class Csv2Ifc:
         unit = self.units.get(symbol, None)
         if unit:
             return unit
+        assert self.file
         unit = self.file.create_entity(
             "IfcContextDependentUnit",
             self.file.create_entity("IfcDimensionalExponents", 0, 0, 0, 0, 0, 0, 0),
diff --git a/src/ifc5d/ifc5d/ifc5Dspreadsheet.py b/src/ifc5d/ifc5d/ifc5Dspreadsheet.py
index c237d3334c..c3a3663b3e 100644
--- a/src/ifc5d/ifc5d/ifc5Dspreadsheet.py
+++ b/src/ifc5d/ifc5d/ifc5Dspreadsheet.py
@@ -61,7 +61,7 @@ class IfcDataGetter:
 
     @staticmethod
     def get_root_costs(cost_schedule: ifcopenshell.entity_instance) -> list[ifcopenshell.entity_instance]:
-        return [obj for rel in cost_schedule.Controls or [] for obj in rel.RelatedObjects or []]
+        return ifcopenshell.util.cost.get_root_cost_items(cost_schedule)
 
     @staticmethod
     def get_cost_item_values(cost_item: Union[ifcopenshell.entity_instance, None]) -> Union[list[dict[str, Any]], None]:
@@ -85,6 +85,9 @@ class IfcDataGetter:
 
     @staticmethod
     def process_categories(cost_item: ifcopenshell.entity_instance, categories: set[str]) -> set[str]:
+        """
+        :param categories: A set to fill with categories.
+        """
         for cost_value in cost_item.CostValues or []:
             if cost_value.Category:
                 categories.add("{}{}".format(cost_value.Category, " Cost"))
@@ -93,14 +96,16 @@ class IfcDataGetter:
     @staticmethod
     def process_cost_item_categories(cost_item: ifcopenshell.entity_instance, categories: set[str]) -> set[str]:
         IfcDataGetter.process_categories(cost_item, categories)
-        for rel in cost_item.IsNestedBy or []:
-            for child in rel.RelatedObjects or []:
-                IfcDataGetter.process_cost_item_categories(child, categories)
+        for child in ifcopenshell.util.cost.get_nested_cost_items(cost_item):
+            IfcDataGetter.process_cost_item_categories(child, categories)
         return categories
 
     @staticmethod
     def get_cost_rates_categories(schedule: ifcopenshell.entity_instance) -> set[str]:
-        categories = set()
+        """
+        :param categories: A set to fill with categories.
+        """
+        categories: set[str] = set()
         for cost_item in IfcDataGetter.get_root_costs(schedule):
             IfcDataGetter.process_cost_item_categories(cost_item, categories)
         return categories
@@ -110,16 +115,14 @@ class IfcDataGetter:
         file: ifcopenshell.file,
         cost_item: ifcopenshell.entity_instance,
         cost_items_data: list[CostItem],
-        index: int,
+        index: int = 1,
         hierarchy: str = "1",
     ) -> None:
         """
         :param cost_items_data: A list to fill with cost items.
+        :param index: Current hierarchy depth.
         """
 
-        def listToString(s):
-            return ", ".join([str(i) for i in s])
-
         quantity_data = IfcDataGetter.get_cost_item_quantity(file, cost_item)
         cost_values_data = IfcDataGetter.get_cost_item_values(cost_item)
 
@@ -150,15 +153,13 @@ class IfcDataGetter:
 
         index += 1
         child_hierarchy = hierarchy + ".1"
-        for nested_cost in [obj for rel in cost_item.IsNestedBy or [] for obj in rel.RelatedObjects or []]:
+        for i, nested_cost in enumerate(ifcopenshell.util.cost.get_nested_cost_items(cost_item), 1):
+            child_hierarchy = f"{hierarchy}.{i}"
             IfcDataGetter.process_cost_data(file, nested_cost, cost_items_data, index, child_hierarchy)
-            child_hierarchy = (
-                ".".join(child_hierarchy.split(".")[:-1]) + "." + str(int(child_hierarchy.split(".")[-1]) + 1)
-            )
 
     @staticmethod
     def get_cost_items_data(file: ifcopenshell.file, schedule: ifcopenshell.entity_instance) -> list[CostItem]:
-        cost_items_data: list[cost_item] = []
+        cost_items_data: list[CostItem] = []
         index = 0
         for cost_item in IfcDataGetter.get_root_costs(schedule):
             IfcDataGetter.process_cost_data(file, cost_item, cost_items_data, index)
@@ -278,6 +279,8 @@ class Ifc5Dwriter:
         self.colours = self.default_colors.copy()
 
     def parse(self):
+        """Fill ``sheet_data`` from ``cost_schedules``."""
+        self.sheet_data = {}
         counter: Counter[str] = Counter()
         for cost_schedule in self.cost_schedules:
             sheet_id = cost_schedule.id()
@@ -326,7 +329,6 @@ class Ifc5Dwriter:
 
     def write(self):
         self.cost_schedules = IfcDataGetter.get_schedules(self.file, self.cost_schedule)
-        self.sheet_data = {}
         self.parse()