From 9ec511493ac3edf4771a7413e9f6639229767ae7 Mon Sep 17 00:00:00 2001
From: Adam Johnson <AdamJohnso@gmail.com>
Date: Wed, 7 Jun 2017 15:48:53 -0400
Subject: [PATCH] Update Logic Modifiers to ID props

---
 korman/properties/modifiers/logic.py | 29 ++++++++++++++++------------
 korman/ui/modifiers/logic.py         |  2 +-
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/korman/properties/modifiers/logic.py b/korman/properties/modifiers/logic.py
index 340bf15..339fd9e 100644
--- a/korman/properties/modifiers/logic.py
+++ b/korman/properties/modifiers/logic.py
@@ -19,29 +19,31 @@ from PyHSPlasma import *
 
 from .base import PlasmaModifierProperties
 from ...exporter import ExportError
+from ... import idprops
 
 game_versions = [("pvPrime", "Ages Beyond Myst (63.11)", "Targets the original Uru (Live) game"),
                  ("pvPots", "Path of the Shell (63.12)", "Targets the most recent offline expansion pack"),
                  ("pvMoul", "Myst Online: Uru Live (70)", "Targets the most recent online game")]
 
-class PlasmaVersionedNodeTree(bpy.types.PropertyGroup):
+class PlasmaVersionedNodeTree(idprops.IDPropMixin, bpy.types.PropertyGroup):
     name = StringProperty(name="Name")
     version = EnumProperty(name="Version",
                            description="Plasma versions this node tree exports under",
                            items=game_versions,
                            options={"ENUM_FLAG"},
                            default=set(list(zip(*game_versions))[0]))
-    node_tree_name = StringProperty(name="Node Tree",
-                                    description="Node Tree to export")
+    node_tree = PointerProperty(name="Node Tree",
+                                description="Node Tree to export",
+                                type=bpy.types.NodeTree)
     node_name = StringProperty(name="Node Ref",
                                description="Attach a reference to this node")
 
-    @property
-    def node_tree(self):
-        try:
-            return bpy.data.node_groups[self.node_tree_name]
-        except KeyError:
-            raise ExportError("Node Tree '{}' does not exist!".format(self.node_tree_name))
+    @classmethod
+    def _idprop_mapping(cls):
+        return {"node_tree": "node_tree_name"}
+
+    def _idprop_sources(self):
+        return {"node_tree_name": bpy.data.node_groups}
 
 
 class PlasmaAdvancedLogic(PlasmaModifierProperties):
@@ -60,19 +62,22 @@ class PlasmaAdvancedLogic(PlasmaModifierProperties):
         for i in self.logic_groups:
             our_versions = [globals()[j] for j in i.version]
             if version in our_versions:
+                if i.node_tree is None:
+                    raise ExportError("'{}': Advanced Logic is missing a node tree for '{}'".format(bo.name, i.version))
+
                 # If node_name is defined, then we're only adding a reference. We will make sure that
                 # the entire node tree is exported once before the post_export step, however.
                 if i.node_name:
-                    exporter.want_node_trees[i.node_tree_name] = (bo, so)
+                    exporter.want_node_trees[i.node_tree.name] = (bo, so)
                     node = i.node_tree.nodes.get(i.node_name, None)
                     if node is None:
-                        raise ExportError("Node '{}' does not exist in '{}'".format(i.node_name, i.node_tree_name))
+                        raise ExportError("Node '{}' does not exist in '{}'".format(i.node_name, i.node_tree.name))
                     # We are going to assume get_key will do the adding correctly. Single modifiers
                     # should fetch the appropriate SceneObject before doing anything, so this will
                     # be a no-op in that case. Multi modifiers should accept any SceneObject, however
                     node.get_key(exporter, so)
                 else:
-                    exporter.node_trees_exported.add(i.node_tree_name)
+                    exporter.node_trees_exported.add(i.node_tree.name)
                     i.node_tree.export(exporter, bo, so)
 
     def harvest_actors(self):
diff --git a/korman/ui/modifiers/logic.py b/korman/ui/modifiers/logic.py
index 9dc27f1..c96f675 100644
--- a/korman/ui/modifiers/logic.py
+++ b/korman/ui/modifiers/logic.py
@@ -39,7 +39,7 @@ def advanced_logic(modifier, layout, context):
     if modifier.logic_groups:
         logic = modifier.logic_groups[modifier.active_group_index]
         layout.row().prop_menu_enum(logic, "version")
-        layout.prop_search(logic, "node_tree_name", bpy.data, "node_groups", icon="NODETREE")
+        layout.prop(logic, "node_tree", icon="NODETREE")
         try:
             layout.prop_search(logic, "node_name", logic.node_tree, "nodes", icon="NODE")
         except: