From e9ee49c8a677b4a5121ebe0aeb453b59ebcae4f1 Mon Sep 17 00:00:00 2001
From: Adam Johnson <AdamJohnso@gmail.com>
Date: Wed, 27 Feb 2019 22:53:22 -0500
Subject: [PATCH] Fix #138

Due to memory management issues, it is likely impossible to remove a
node in its own init() callback without crashing Blender. Now, we will
deactivate any output node operators if an output node is already
present in the tree.
---
 korman/nodes/__init__.py  | 36 +++++++++++++++++++++++++++++-------
 korman/nodes/node_core.py | 13 +++++--------
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/korman/nodes/__init__.py b/korman/nodes/__init__.py
index 28c65eb..57aa758 100644
--- a/korman/nodes/__init__.py
+++ b/korman/nodes/__init__.py
@@ -50,26 +50,48 @@ _kategory_names = {
     "SV": "Soft Volume",
 }
 
+class PlasmaNodeItem(NodeItem):
+    def __init__(self, **kwargs):
+        self._poll_add = kwargs.pop("poll_add", None)
+        super().__init__(**kwargs)
+
+    @staticmethod
+    def draw(self, layout, context):
+        # Blender's NodeItem completely unlists anything that polls False. We would rather
+        # display the item but disabled so the user has a hint that it exists but cannot be
+        # used at the present time.
+        if self._poll_add is not None and not self._poll_add(context):
+            row = layout.row()
+            row.enabled = False
+            NodeItem.draw(self, row, context)
+        else:
+            NodeItem.draw(self, layout, context)
+
+
 # Now, generate the categories as best we can...
 _kategories = {}
 for cls in dict(globals()).values():
     if inspect.isclass(cls):
+        item = {}
         if not issubclass(cls, PlasmaNodeBase):
             continue
         if not issubclass(cls, bpy.types.Node):
             continue
         if issubclass(cls, PlasmaDeprecatedNode):
             continue
-    else:
-        continue
-    try:
-        _kategories[cls.bl_category].append(cls)
-    except LookupError:
-        _kategories[cls.bl_category] = [cls,]
+        if issubclass(cls, PlasmaTreeOutputNodeBase):
+            item["poll_add"] = PlasmaTreeOutputNodeBase.poll_add
+
+        item["nodetype"] = cls.bl_idname
+        item["label"] = cls.bl_label
+
+        kat = _kategories.setdefault(cls.bl_category, [])
+        kat.append(item)
+
 _actual_kategories = []
 for i in sorted(_kategories.keys(), key=lambda x: _kategory_names[x]):
     # Note that even though we're sorting the category names, Blender appears to not care...
-    _kat_items = [NodeItem(j.bl_idname) for j in sorted(_kategories[i], key=lambda x: x.bl_label)]
+    _kat_items = [PlasmaNodeItem(**j) for j in sorted(_kategories[i], key=lambda x: x["label"])]
     _actual_kategories.append(PlasmaNodeCategory(i, _kategory_names[i], items=_kat_items))
 
 def register():
diff --git a/korman/nodes/node_core.py b/korman/nodes/node_core.py
index 9710b53..74bbd04 100644
--- a/korman/nodes/node_core.py
+++ b/korman/nodes/node_core.py
@@ -313,14 +313,11 @@ class PlasmaNodeBase:
 
 class PlasmaTreeOutputNodeBase(PlasmaNodeBase):
     """Represents the final output of a node tree"""
-    def init(self, context):
-        nodes = self.id_data.nodes
-
-        # There can only be one of these nodes per tree, so let's make sure I'm the only one.
-        for i in nodes:
-            if isinstance(i, self.__class__) and i != self:
-                nodes.remove(self)
-                return
+    @classmethod
+    def poll_add(cls, context):
+        # There can only be one of these nodes per tree, so we will only allow this to be
+        # added if no other output nodes are found.
+        return not any((isinstance(node, cls) for node in context.space_data.node_tree.nodes))
 
 
 class PlasmaNodeSocketBase: