From 53d64348fd72bddb6f168fb3e18aa76263dd8195 Mon Sep 17 00:00:00 2001
From: Dion Moult <dion@thinkmoult.com>
Date: Sat, 7 Dec 2019 21:33:44 +1100
Subject: [PATCH] Allow people to configure exportable contexts

---
 src/ifcblenderexport/blenderbim/__init__.py   |  2 ++
 src/ifcblenderexport/blenderbim/export_ifc.py | 27 +------------------
 src/ifcblenderexport/blenderbim/operator.py   | 13 +++++++++
 src/ifcblenderexport/blenderbim/prop.py       | 11 ++++++++
 4 files changed, 27 insertions(+), 26 deletions(-)

diff --git a/src/ifcblenderexport/blenderbim/__init__.py b/src/ifcblenderexport/blenderbim/__init__.py
index 207bed474b..dea12bc80a 100644
--- a/src/ifcblenderexport/blenderbim/__init__.py
+++ b/src/ifcblenderexport/blenderbim/__init__.py
@@ -98,6 +98,7 @@ def menu_func_import(self, context):
 def register():
     for cls in classes:
         bpy.utils.register_class(cls)
+    bpy.app.handlers.load_post.append(prop.setDefaultProperties)
     bpy.types.TOPBAR_MT_file_export.append(menu_func_export)
     bpy.types.TOPBAR_MT_file_import.append(menu_func_import)
     bpy.types.Scene.BIMProperties = bpy.props.PointerProperty(type=prop.BIMProperties)
@@ -112,6 +113,7 @@ def register():
 def unregister():
     for cls in reversed(classes):
         bpy.utils.unregister_class(cls)
+    bpy.app.handlers.load_post.remove(prop.setDefaultProperties)
     bpy.types.TOPBAR_MT_file_export.remove(menu_func_export)
     bpy.types.TOPBAR_MT_file_import.remove(menu_func_import)
     del(bpy.types.Scene.BIMProperties)
diff --git a/src/ifcblenderexport/blenderbim/export_ifc.py b/src/ifcblenderexport/blenderbim/export_ifc.py
index b8b3daa731..f4a1b74f11 100644
--- a/src/ifcblenderexport/blenderbim/export_ifc.py
+++ b/src/ifcblenderexport/blenderbim/export_ifc.py
@@ -2019,29 +2019,4 @@ class IfcExportSettings:
                              'SECTION_VIEW', 'ELEVATION_VIEW', 'USERDEFINED', 'NOTDEFINED']
         self.should_export_all_materials_as_styled_items = False
         self.should_use_presentation_style_assignment = False
-        # TODO make this configurable via UI
-        self.context_tree = self.build_context_tree()
-
-    def build_context_tree(self):
-        tree = []
-        for context in self.contexts:
-            subcontexts = []
-            for subcontext in self.subcontexts + self.generated_subcontexts:
-                target_views = []
-                for target_view in self.target_views:
-                    if context == 'Model' \
-                            and target_view != 'MODEL_VIEW':
-                        continue
-                    elif context == 'Plan' \
-                            and target_view == 'MODEL_VIEW':
-                        continue
-                    target_views.append(target_view)
-                subcontexts.append({
-                    'name': subcontext,
-                    'target_views': target_views
-                })
-            tree.append({
-                'name': context,
-                'subcontexts': subcontexts
-            })
-        return tree
+        self.context_tree = []
diff --git a/src/ifcblenderexport/blenderbim/operator.py b/src/ifcblenderexport/blenderbim/operator.py
index 1c1dd29e7a..003ec1d4bd 100644
--- a/src/ifcblenderexport/blenderbim/operator.py
+++ b/src/ifcblenderexport/blenderbim/operator.py
@@ -36,6 +36,19 @@ class ExportIFC(bpy.types.Operator):
         ifc_export_settings.has_representations = bpy.context.scene.BIMProperties.export_has_representations
         ifc_export_settings.should_export_all_materials_as_styled_items = bpy.context.scene.BIMProperties.export_should_export_all_materials_as_styled_items
         ifc_export_settings.should_use_presentation_style_assignment = bpy.context.scene.BIMProperties.export_should_use_presentation_style_assignment
+        ifc_export_settings.context_tree = []
+        for context in ['model', 'plan']:
+            if getattr(bpy.context.scene.BIMProperties, 'has_{}_context'.format(context)):
+                subcontexts = {}
+                for subcontext in getattr(bpy.context.scene.BIMProperties, '{}_subcontexts'.format(context)):
+                    subcontexts.setdefault(subcontext.name, []).append(subcontext.target_view)
+                ifc_export_settings.context_tree.append({
+                    'name': context.title(),
+                    'subcontexts': [
+                        { 'name': key, 'target_views': value }
+                        for key, value in subcontexts.items()
+                    ]
+                })
         ifc_parser = export_ifc.IfcParser(ifc_export_settings)
         ifc_schema = export_ifc.IfcSchema(ifc_export_settings)
         qto_calculator = export_ifc.QtoCalculator()
diff --git a/src/ifcblenderexport/blenderbim/prop.py b/src/ifcblenderexport/blenderbim/prop.py
index 6fcb5b3c61..22f8ea4d33 100644
--- a/src/ifcblenderexport/blenderbim/prop.py
+++ b/src/ifcblenderexport/blenderbim/prop.py
@@ -3,7 +3,9 @@ import os
 import csv
 from pathlib import Path
 from . import export_ifc
+import bpy
 from bpy.types import PropertyGroup
+from bpy.app.handlers import persistent
 from bpy.props import (
     StringProperty,
     EnumProperty,
@@ -34,6 +36,15 @@ contexts_enum = []
 subcontexts_enum = []
 target_views_enum = []
 
+@persistent
+def setDefaultProperties(scene):
+    if bpy.context.scene.BIMProperties.has_model_context \
+            and len(bpy.context.scene.BIMProperties.model_subcontexts) == 0:
+        subcontext = bpy.context.scene.BIMProperties.model_subcontexts.add()
+        subcontext.name = 'Body'
+        subcontext.target_view = 'MODEL_VIEW'
+
+
 def getIfcPredefinedTypes(self, context):
     global types_enum
     if len(types_enum) < 1: