From 6502b2c18da290c9cb5666a9887e5b196276cf81 Mon Sep 17 00:00:00 2001
From: QwiglyDee <qwiglydee@gmail.com>
Date: Tue, 12 Jan 2021 04:02:39 +0700
Subject: [PATCH] little drawing helper

---
 .../blenderbim/bim/decoration.py              |  7 +++----
 src/ifcblenderexport/blenderbim/bim/helper.py | 21 +++++++++++++++++++
 .../blenderbim/bim/operator.py                | 18 ++++++----------
 3 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/src/ifcblenderexport/blenderbim/bim/decoration.py b/src/ifcblenderexport/blenderbim/bim/decoration.py
index b52a3b18db..44b756d9a7 100644
--- a/src/ifcblenderexport/blenderbim/bim/decoration.py
+++ b/src/ifcblenderexport/blenderbim/bim/decoration.py
@@ -13,6 +13,7 @@ import gpu
 import bgl
 from gpu.types import GPUShader, GPUBatch, GPUIndexBuf, GPUVertBuf, GPUVertFormat
 from gpu_extras.batch import batch_for_shader
+from . import helper
 
 
 class BaseDecorator():
@@ -1031,11 +1032,9 @@ class DecorationsHandler():
         self.decorators = [cls() for cls in self.decorators_classes]
 
     def __call__(self, context):
-        props = context.scene.DocProperties
-        if props.active_drawing_index is None or len(props.drawings) == 0:
+        collection, _ = helper.get_active_drawing(context.scene)
+        if collection is None:
             return
-        drawing = props.drawings[props.active_drawing_index]
-        collection = bpy.data.collections.get("IfcGroup/" + drawing.name)
 
         for decorator in self.decorators:
             for obj in decorator.get_objects(collection):
diff --git a/src/ifcblenderexport/blenderbim/bim/helper.py b/src/ifcblenderexport/blenderbim/bim/helper.py
index 185f631121..8308a182ae 100644
--- a/src/ifcblenderexport/blenderbim/bim/helper.py
+++ b/src/ifcblenderexport/blenderbim/bim/helper.py
@@ -299,6 +299,27 @@ def parse_diagram_scale(camera):
     return float(numerator) / float(denominator)
 
 
+def get_project_collection(scene):
+    """Get main project collection"""
+
+    colls = [c for c in scene.collection.children if c.name.startswith('IfcProject')]
+    if len(colls) != 1:
+        raise RuntimeError("project collection missing or not unique")
+    return colls[0]
+
+
+def get_active_drawing(scene):
+    """Get active drawing collection and camera"""
+    props = scene.DocProperties
+    if props.active_drawing_index is None or len(props.drawings) == 0:
+        return None, None
+    try:
+        drawing = props.drawings[props.active_drawing_index]
+        return scene.collection.children['Views'].children[f"IfcGroup/{drawing.name}"], drawing.camera
+    except (KeyError, IndexError):
+        raise RuntimeError("missing drawing collection")
+
+
 def ortho_view_frame(camera, margin=0.015):
     """Calculates 2d bounding box of camera view area.
 
diff --git a/src/ifcblenderexport/blenderbim/bim/operator.py b/src/ifcblenderexport/blenderbim/bim/operator.py
index a7ca2e9c93..2e465cae67 100644
--- a/src/ifcblenderexport/blenderbim/bim/operator.py
+++ b/src/ifcblenderexport/blenderbim/bim/operator.py
@@ -3757,24 +3757,18 @@ class CopyGrid(bpy.types.Operator):
     bl_options = {"UNDO"}
 
     def execute(self, context):
-        props = context.scene.DocProperties
-        if props.active_drawing_index is None or len(props.drawings) == 0:
-            return {"CANCELLED"}
-        drawing = props.drawings[props.active_drawing_index]
-        collection = bpy.data.collections.get("IfcGroup/" + drawing.name)
+        proj_coll = helper.get_project_collection(context.scene)
+        view_coll, camera = helper.get_active_drawing(context.scene)
 
-        existing = [obj for obj in collection.objects if obj.name.startswith("IfcGridAxis")]
+        existing = [obj for obj in view_coll.objects if obj.name.startswith("IfcGridAxis")]
         for obj in existing:
-            collection.objects.unlink(obj)
+            view_coll.objects.unlink(obj)
 
         source = [
             obj
-            for coll in bpy.data.collections
-            if coll.name.startswith("IfcGrid")
-            for obj in coll.all_objects
+            for obj in proj_coll.all_objects
             if obj.name.startswith("IfcGridAxis")
         ]
-        camera = [obj for obj in collection.all_objects if obj.type == "CAMERA"][0]
 
         clipping = camera.data.type == "ORTHO"
         bounds = helper.ortho_view_frame(camera.data) if clipping else None
@@ -3804,6 +3798,6 @@ class CopyGrid(bpy.types.Operator):
             dst = src.copy()
             dst.data = bpy.data.meshes.new(dst.name)
             bm.to_mesh(dst.data)
-            collection.objects.link(dst)
+            view_coll.objects.link(dst)
 
         return {"FINISHED"}