some refactor for bbim drawings svg serialization

Made some changes to svg serialization code by moving chunks of code so `.generate_linework()` will be easier to read
- moved gathering linework contexts to a separate method
- moved common code for elements serialization to a function
- replaced self.yield_from_iterator(it) with it.__iter__

Functionally code is exactly the same.

src/blenderbim/blenderbim/bim/module/drawing/operator.py

@@ -44,6 +44,7 @@ import blenderbim.bim.module.drawing.helper as helper
 import blenderbim.bim.export_ifc
 from blenderbim.bim.module.drawing.decoration import CutDecorator
 from blenderbim.bim.module.drawing.data import DecoratorData, DrawingsData
+from typing import NamedTuple, List
 from lxml import etree
 from mathutils import Vector, Color, Matrix
 from timeit import default_timer as timer
@@ -77,6 +78,11 @@ class Operator:
         return {"FINISHED"}
 
 
+class LineworkContexts(NamedTuple):
+    body: List[List[int]]
+    annotation: List[List[int]]
+
+
 class AddAnnotationType(bpy.types.Operator, Operator):
     bl_idname = "bim.add_annotation_type"
     bl_label = "Add Annotation Type"
@@ -367,48 +373,7 @@ class CreateDrawing(bpy.types.Operator):
         self.svg_writer.create_blank_svg(svg_path).draw_underlay(context.scene.render.filepath).save()
         return svg_path
 
-    def generate_linework(self, context):
+    def get_linework_contexts(self, ifc, target_view) -> LineworkContexts:
-        if not ifcopenshell.util.element.get_pset(self.drawing, "EPset_Drawing", "HasLinework"):
-            return
-        svg_path = self.get_svg_path(cache_type="linework")
-        if os.path.isfile(svg_path) and self.props.should_use_linework_cache:
-            return svg_path
-
-        with profile("sync"):
-            # All very hackish whilst prototyping
-            exporter = blenderbim.bim.export_ifc.IfcExporter(None)
-            exporter.file = tool.Ifc.get()
-            invalidated_elements = exporter.sync_all_objects()
-            invalidated_elements += exporter.sync_edited_objects()
-            invalidated_guids = [e.GlobalId for e in invalidated_elements if hasattr(e, "GlobalId")]
-
-        # If we have already calculated it in the SVG in the past, don't recalculate
-        edited_guids = set()
-        for obj in IfcStore.edited_objs:
-            element = tool.Ifc.get_entity(obj)
-            edited_guids.add(element.GlobalId) if hasattr(element, "GlobalId") else None
-        cached_linework = set()
-        if os.path.isfile(svg_path):
-            tree = etree.parse(svg_path)
-            root = tree.getroot()
-            cached_linework = {
-                el.get("{http://www.ifcopenshell.org/ns}guid")
-                for el in root.findall(".//{http://www.w3.org/2000/svg}g[@{http://www.ifcopenshell.org/ns}guid]")
-            }
-        cached_linework -= edited_guids
-
-        files = {context.scene.BIMProperties.ifc_file: tool.Ifc.get()}
-
-        for ifc_path, ifc in files.items():
-            # Don't use draw.main() just whilst we're prototyping and experimenting
-            ifc_hash = hashlib.md5(ifc_path.encode("utf-8")).hexdigest()
-            ifc_cache_path = os.path.join(context.scene.BIMProperties.data_dir, "cache", f"{ifc_hash}.h5")
-
-            # Get all representation contexts to see what we're dealing with.
-            # Drawings only draw bodies and annotations (and facetation, due to a Revit bug).
-            # A drawing prioritises a target view context first, followed by a model view context as a fallback.
-            # Specifically for PLAN_VIEW and REFLECTED_PLAN_VIEW, any Plan context is also prioritised.
-
         plan_body_target_contexts = []
         plan_body_model_contexts = []
         model_body_target_contexts = []
@@ -419,8 +384,6 @@ class CreateDrawing(bpy.types.Operator):
         model_annotation_target_contexts = []
         model_annotation_model_contexts = []
 
-            target_view = ifcopenshell.util.element.get_psets(self.camera_element)["EPset_Drawing"]["TargetView"]
-
         for rep_context in ifc.by_type("IfcGeometricRepresentationContext"):
             if rep_context.is_a("IfcGeometricRepresentationSubContext"):
                 if rep_context.ContextType == "Plan":
@@ -479,6 +442,79 @@ class CreateDrawing(bpy.types.Operator):
             ]
         )
 
+        return LineworkContexts(body_contexts, annotation_contexts)
+
+    def serialize_contexts_elements(
+        self, ifc, tree: ifcopenshell.geom.tree, contexts: LineworkContexts, context_type, drawing_elements, target_view
+    ):
+        drawing_elements = drawing_elements.copy()
+        contexts = getattr(contexts, context_type)
+        for context in contexts:
+            with profile(f"Processing {context_type} context"):
+                if not context or not drawing_elements:
+                    continue
+                geom_settings = ifcopenshell.geom.settings(
+                    DISABLE_TRIANGULATION=True, STRICT_TOLERANCE=True, INCLUDE_CURVES=True
+                )
+                if ifc.by_id(context[0]).ContextType == "Plan" and "PLAN_VIEW" in target_view:
+                    offset = ifcopenshell.ifcopenshell_wrapper.float_array_3()
+                    # A 2mm Z offset to combat Z-fighting in plan or RCPs
+                    offset[2] = 0.002 if target_view == "PLAN_VIEW" else -0.002
+                    geom_settings.offset = offset
+                geom_settings.set_context_ids(context)
+                it = ifcopenshell.geom.iterator(
+                    geom_settings, ifc, multiprocessing.cpu_count(), include=drawing_elements
+                )
+                processed = set()
+                for elem in it:
+                    processed.add(ifc.by_id(elem.id))
+                    self.serialiser.write(elem)
+                    tree.add_element(elem)
+                drawing_elements -= processed
+
+    def generate_linework(self, context):
+        if not ifcopenshell.util.element.get_pset(self.drawing, "EPset_Drawing", "HasLinework"):
+            return
+        svg_path = self.get_svg_path(cache_type="linework")
+        if os.path.isfile(svg_path) and self.props.should_use_linework_cache:
+            return svg_path
+
+        with profile("sync"):
+            # All very hackish whilst prototyping
+            exporter = blenderbim.bim.export_ifc.IfcExporter(None)
+            exporter.file = tool.Ifc.get()
+            invalidated_elements = exporter.sync_all_objects()
+            invalidated_elements += exporter.sync_edited_objects()
+            invalidated_guids = [e.GlobalId for e in invalidated_elements if hasattr(e, "GlobalId")]
+
+        # If we have already calculated it in the SVG in the past, don't recalculate
+        edited_guids = set()
+        for obj in IfcStore.edited_objs:
+            element = tool.Ifc.get_entity(obj)
+            edited_guids.add(element.GlobalId) if hasattr(element, "GlobalId") else None
+        cached_linework = set()
+        if os.path.isfile(svg_path):
+            tree = etree.parse(svg_path)
+            root = tree.getroot()
+            cached_linework = {
+                el.get("{http://www.ifcopenshell.org/ns}guid")
+                for el in root.findall(".//{http://www.w3.org/2000/svg}g[@{http://www.ifcopenshell.org/ns}guid]")
+            }
+        cached_linework -= edited_guids
+
+        files = {context.scene.BIMProperties.ifc_file: tool.Ifc.get()}
+
+        for ifc_path, ifc in files.items():
+            # Don't use draw.main() just whilst we're prototyping and experimenting
+            ifc_hash = hashlib.md5(ifc_path.encode("utf-8")).hexdigest()
+            ifc_cache_path = os.path.join(context.scene.BIMProperties.data_dir, "cache", f"{ifc_hash}.h5")
+
+            # Get all representation contexts to see what we're dealing with.
+            # Drawings only draw bodies and annotations (and facetation, due to a Revit bug).
+            # A drawing prioritises a target view context first, followed by a model view context as a fallback.
+            # Specifically for PLAN_VIEW and REFLECTED_PLAN_VIEW, any Plan context is also prioritised.
+            target_view = ifcopenshell.util.element.get_psets(self.camera_element)["EPset_Drawing"]["TargetView"]
+            contexts = self.get_linework_contexts(ifc, target_view)
             drawing_elements = tool.Drawing.get_drawing_elements(self.camera_element)
 
             self.setup_serialiser(ifc, target_view)
@@ -487,60 +523,15 @@ class CreateDrawing(bpy.types.Operator):
             tree = ifcopenshell.geom.tree()
             tree.enable_face_styles(True)
 
-            elements = drawing_elements.copy()
+            self.serialize_contexts_elements(ifc, tree, contexts, "body", drawing_elements, target_view)
-            for body_context in body_contexts:
+            self.serialize_contexts_elements(ifc, tree, contexts, "annotation", drawing_elements, target_view)
-                with profile("Processing body context"):
-                    if body_context and elements:
-                        geom_settings = ifcopenshell.geom.settings(
-                            DISABLE_TRIANGULATION=True, STRICT_TOLERANCE=True, INCLUDE_CURVES=True
-                        )
-                        if ifc.by_id(body_context[0]).ContextType == "Plan" and "PLAN_VIEW" in target_view:
-                            offset = ifcopenshell.ifcopenshell_wrapper.float_array_3()
-                            # A 2mm Z offset to combat Z-fighting in plan or RCPs
-                            offset[2] = 0.002 if target_view == "PLAN_VIEW" else -0.002
-                            geom_settings.offset = offset
-                        geom_settings.set_context_ids(body_context)
-                        it = ifcopenshell.geom.iterator(
-                            geom_settings, ifc, multiprocessing.cpu_count(), include=elements
-                        )
-                        it.set_cache(cache)
-                        processed = set()
-                        for elem in self.yield_from_iterator(it):
-                            processed.add(ifc.by_id(elem.id))
-                            self.serialiser.write(elem)
-                            tree.add_element(elem)
-                        elements -= processed
-
-            elements = drawing_elements.copy()
-            for annotation_context in annotation_contexts:
-                with profile("Processing annotation context"):
-                    if annotation_context and elements:
-                        geom_settings = ifcopenshell.geom.settings(
-                            DISABLE_TRIANGULATION=True, STRICT_TOLERANCE=True, INCLUDE_CURVES=True
-                        )
-                        if ifc.by_id(annotation_context[0]).ContextType == "Plan" and "PLAN_VIEW" in target_view:
-                            offset = ifcopenshell.ifcopenshell_wrapper.float_array_3()
-                            # A 2mm Z offset to combat Z-fighting in plan or RCPs
-                            offset[2] = 0.002 if target_view == "PLAN_VIEW" else -0.002
-                            geom_settings.offset = offset
-                        geom_settings.set_context_ids(annotation_context)
-                        it = ifcopenshell.geom.iterator(
-                            geom_settings, ifc, multiprocessing.cpu_count(), include=elements
-                        )
-                        it.set_cache(cache)
-                        processed = set()
-                        for elem in self.yield_from_iterator(it):
-                            processed.add(ifc.by_id(elem.id))
-                            self.serialiser.write(elem)
-                            tree.add_element(elem)
-                        elements -= processed
 
             if self.camera_element not in drawing_elements:
                 with profile("Camera element"):
                     # The camera must always be included, regardless of any include/exclude filters.
                     geom_settings = ifcopenshell.geom.settings(DISABLE_TRIANGULATION=True, STRICT_TOLERANCE=True)
                     it = ifcopenshell.geom.iterator(geom_settings, ifc, include=[self.camera_element])
-                    for elem in self.yield_from_iterator(it):
+                    for elem in it:
                         self.serialiser.write(elem)
 
         with profile("Finalizing"):
@@ -884,13 +875,6 @@ class CreateDrawing(bpy.types.Operator):
         # the style of the face when running tree.select_ray()
         # tree.enable_face_styles(True)
 
-    def yield_from_iterator(self, it):
-        if it.initialize():
-            while True:
-                yield it.get()
-                if not it.next():
-                    break
-
     def get_svg_classes(self, element):
         classes = [element.is_a()]
         material = ifcopenshell.util.element.get_material(element, should_skip_usage=True)