Nicer registration mechanism for node categories. The lists of node categories and items are now stored in a dictionary with an identifier key, so they can be registered and unregistered individually. The Add menu is now persistent and gets extended with a draw function for each of the registered node category lists.

This allows pynodes to define their own list of node categories and items and register it at runtime without interfering with the standard nodes.

3 files changed, 67 insertions, 33 deletions

diff --git a/release/scripts/modules/nodeitems_utils.py b/release/scripts/modules/nodeitems_utils.py
index 522e61769ea..b3bd7cf3151 100644
--- a/release/scripts/modules/nodeitems_utils.py
+++ b/release/scripts/modules/nodeitems_utils.py
@@ -21,9 +21,6 @@ import bpy
 from bpy.types import Menu, Panel
 
 
-node_categories = []
-
-
 class NodeCategory():
     @classmethod
     def poll(cls, context):
@@ -60,12 +57,13 @@ class NodeItem():
             return getattr(bpy.types, self.nodetype).bl_rna.name
 
 
-# Empty base class to detect subclasses in bpy.types
-class NodeCategoryUI():
-    pass
+_node_categories = {}
 
+def register_node_categories(identifier, cat_list):
+    if identifier in _node_categories:
+        raise KeyError("Node categories list '%s' already registered" % identifier)
+        return
 
-def register_node_ui():
     # works as draw function for both menus and panels
     def draw_node_item(self, context):
         layout = self.layout
@@ -80,15 +78,17 @@ def register_node_ui():
                 ops.name = setting[0]
                 ops.value = setting[1]
 
-    for cat in node_categories:
-        menu = type("NODE_MT_category_"+cat.identifier, (bpy.types.Menu, NodeCategoryUI), {
+    menu_types = []
+    panel_types = []
+    for cat in cat_list:
+        menu_type = type("NODE_MT_category_"+cat.identifier, (bpy.types.Menu,), {
             "bl_space_type" : 'NODE_EDITOR',
             "bl_label" : cat.name,
             "category" : cat,
             "poll" : cat.poll,
             "draw" : draw_node_item,
             })
-        panel = type("NODE_PT_category_"+cat.identifier, (bpy.types.Panel, NodeCategoryUI), {
+        panel_type = type("NODE_PT_category_"+cat.identifier, (bpy.types.Panel,), {
             "bl_space_type" : 'NODE_EDITOR',
             "bl_region_type" : 'TOOLS',
             "bl_label" : cat.name,
@@ -97,32 +97,44 @@ def register_node_ui():
             "poll" : cat.poll,
             "draw" : draw_node_item,
             })
-        bpy.utils.register_class(menu)
-        bpy.utils.register_class(panel)
 
+        menu_types.append(menu_type)
+        panel_types.append(panel_type)
+
+        bpy.utils.register_class(menu_type)
+        bpy.utils.register_class(panel_type)
 
     def draw_add_menu(self, context):
         layout = self.layout
 
-        layout.operator_context = 'INVOKE_DEFAULT'
-        op = layout.operator("node.add_search", text="Search ...")
-
-        for cat in node_categories:
+        for cat in cat_list:
             if cat.poll(context):
                 layout.menu("NODE_MT_category_%s" % cat.identifier)
 
-    add_menu = type("NODE_MT_add", (bpy.types.Menu, NodeCategoryUI), {
-        "bl_space_type" : 'NODE_EDITOR',
-        "bl_label" : "Add",
-        "draw" : draw_add_menu,
-        })
-    bpy.utils.register_class(add_menu)
+    bpy.types.NODE_MT_add.append(draw_add_menu)
+
+    # stores: (categories list, menu draw function, submenu types, panel types)
+    _node_categories[identifier] = (cat_list, draw_add_menu, menu_types, panel_types)
+
+
+def unregister_node_cat_types(cats):
+    bpy.types.NODE_MT_add.remove(cats[1])
+    for mt in cats[2]:
+        bpy.utils.unregister_class(mt)
+    for pt in cats[3]:
+        bpy.utils.unregister_class(pt)
 
 
-def unregister_node_ui():
+def unregister_node_categories(identifier=None):
     # unregister existing UI classes
-    for c in NodeCategoryUI.__subclasses__():
-        if hasattr(c, "bl_rna"):
-            bpy.utils.unregister_class(c)
-            del c
+    if identifier:
+        cat_types = _node_categories.get(identifier, None)
+        if cat_types:
+            unregister_node_cat_types(cat_types)
+        del _node_categories[identifier]
+
+    else:
+        for cat_types in _node_categories.values():
+            unregister_node_cat_types(cat_types)
+        _node_categories.clear()
 
diff --git a/release/scripts/startup/bl_ui/space_node.py b/release/scripts/startup/bl_ui/space_node.py
index c8e11adbdd2..86f25fd3995 100644
--- a/release/scripts/startup/bl_ui/space_node.py
+++ b/release/scripts/startup/bl_ui/space_node.py
@@ -108,6 +108,19 @@ class NODE_HT_header(Header):
         layout.template_running_jobs()
 
 
+class NODE_MT_add(bpy.types.Menu):
+    bl_space_type = 'NODE_EDITOR'
+    bl_label = "Add"
+
+    def draw(self, context):
+        layout = self.layout
+
+        layout.operator_context = 'INVOKE_DEFAULT'
+        op = layout.operator("node.add_search", text="Search ...")
+
+        # actual node submenus are added by draw functions from node categories
+
+
 class NODE_MT_view(Menu):
     bl_label = "View"
 
@@ -316,5 +329,9 @@ class NODE_UL_interface_sockets(bpy.types.UIList):
             layout.template_node_socket(color)
 
 
+def node_draw_tree_view(layout, context):
+    pass
+
+
 if __name__ == "__main__":  # only for live edit.
     bpy.utils.register_module(__name__)
diff --git a/release/scripts/startup/nodeitems_builtins.py b/release/scripts/startup/nodeitems_builtins.py
index 68928bdf2fb..f5945252c11 100644
--- a/release/scripts/startup/nodeitems_builtins.py
+++ b/release/scripts/startup/nodeitems_builtins.py
@@ -63,7 +63,7 @@ def texture_node_group_items(self):
 
 # All standard node categories currently used in nodes.
 
-std_node_categories = [
+shader_node_categories = [
     # Shader Nodes
     ShaderOldNodeCategory("SH_INPUT", "Input", items=[
         NodeItem("ShaderNodeMaterial"),
@@ -184,7 +184,9 @@ std_node_categories = [
     ShaderNewNodeCategory("SH_NEW_GROUP", "Group", items=shader_node_group_items),
     ShaderNewNodeCategory("SH_NEW_LAYOUT", "Layout", items=[
         ]),
+    ]
 
+compositor_node_categories = [
      # Compositor Nodes
      CompositorNodeCategory("CMP_INPUT", "Input", items = [
         NodeItem("CompositorNodeRLayers"),
@@ -287,7 +289,9 @@ std_node_categories = [
     CompositorNodeCategory("CMP_LAYOUT", "Layout", items = [
         NodeItem("CompositorNodeSwitch"),
         ]),
+    ]
 
+texture_node_categories = [
     # Texture Nodes
     TextureNodeCategory("TEX_INPUT", "Input", items = [
         NodeItem("TextureNodeCurveTime"),
@@ -342,14 +346,15 @@ std_node_categories = [
 
 
 def register():
-    # XXX can be made a lot nicer, just get it working for now
-    nodeitems_utils.node_categories = std_node_categories
-    nodeitems_utils.register_node_ui()
+    nodeitems_utils.register_node_categories("SHADER", shader_node_categories)
+    nodeitems_utils.register_node_categories("COMPOSITING", compositor_node_categories)
+    nodeitems_utils.register_node_categories("TEXTURE", texture_node_categories)
 
 
 def unregister():
-    nodeitems_utils.unregister_node_ui()
-    nodeitems_utils.node_categories = []
+    nodeitems_utils.unregister_node_categories("SHADER")
+    nodeitems_utils.unregister_node_categories("COMPOSITING")
+    nodeitems_utils.unregister_node_categories("TEXTURE")
 
 
 if __name__ == "__main__":