Upgrade for the add_search node operator. This now uses the same basic system as the regular add_node operator, with enum items generated from the common node categories system (nodeitems_utils module). This means that any node listed in the regular node Add menu can now also be added via searching, including node groups and the like. The search operator also uses the subsequent transform to make insertion a bit more streamlined.

release/scripts/startup/bl_operators/node.py

@@ -18,7 +18,7 @@
 
 # <pep8-80 compliant>
 
-import bpy
+import bpy, nodeitems_utils
 from bpy.types import Operator, PropertyGroup
 from bpy.props import BoolProperty, CollectionProperty, EnumProperty, IntProperty, StringProperty
 
@@ -62,15 +62,20 @@ class NodeAddOperator():
         else:
             space.cursor_location = tree.view_center
 
-    def create_node(self, context):
+    # XXX explicit node_type argument is usually not necessary, but required to make search operator work:
+    # add_search has to override the 'type' property since it's hardcoded in bpy_operator_wrap.c ...
+    def create_node(self, context, node_type=None):
         space = context.space_data
         tree = space.edit_tree
 
+        if node_type is None:
+            node_type = self.type
+
         # select only the new node
         for n in tree.nodes:
             n.select = False
 
-        node = tree.nodes.new(type=self.type)
+        node = tree.nodes.new(type=node_type)
 
         for setting in self.settings:
             # XXX catch exceptions here?
@@ -151,83 +156,53 @@ class NODE_OT_add_and_link_node(NodeAddOperator, Operator):
         return {'FINISHED'}
 
 
-def node_classes_iter(base=bpy.types.Node):
-    """
-    Yields all true node classes by checking for the is_registered_node_type classmethod.
-    Node types can use specialized subtypes of bpy.types.Node, which are not usable
-    nodes themselves (e.g. CompositorNode).
-    """
-    if base.is_registered_node_type():
-        yield base
-    for subclass in base.__subclasses__():
-        for node_class in node_classes_iter(subclass):
-            yield node_class
-
-
-def node_class_items_iter(node_class, context):
-    identifier = node_class.bl_rna.identifier
-    # XXX Checking for explicit group node types is stupid.
-    # This should be replaced by a generic system of generating
-    # node items via callback.
-    # Group node_tree pointer should also use a poll function to filter the library list,
-    # but cannot do that without a node instance here. A node callback could just use the internal poll function.
-    if identifier in {'ShaderNodeGroup', 'CompositorNodeGroup', 'TextureNodeGroup'}:
-        tree_idname = context.space_data.edit_tree.bl_idname
-        for group in bpy.data.node_groups:
-            if group.bl_idname == tree_idname:
-                # XXX empty string should be replaced by description from tree
-                yield (group.name, "", {"node_tree": group})
-    else:
-        yield (node_class.bl_rna.name, node_class.bl_rna.description, {})
-
-
-def node_items_iter(context):
-    snode = context.space_data
-    if not snode:
-        return
-    tree = snode.edit_tree
-    if not tree:
-        return
-
-    for node_class in node_classes_iter():
-        if node_class.poll(tree):
-            for item in node_class_items_iter(node_class, context):
-                yield (node_class,) + item
-
-
-# Create an enum list from node class items
-def node_type_items_cb(self, context):
-    # XXX Python has to keep a ref to those strings, else they may be freed :(
-    NODE_OT_add_search._enum_str_store = [(str(index), item[1], item[2])
-                                          for index, item in enumerate(node_items_iter(context))]
-    return NODE_OT_add_search._enum_str_store
-
-
 class NODE_OT_add_search(NodeAddOperator, Operator):
     '''Add a node to the active tree'''
     bl_idname = "node.add_search"
     bl_label = "Search and Add Node"
     bl_options = {'REGISTER', 'UNDO'}
 
-    # XXX Python has to keep a ref to the data (strings) generated by enum's callback, else they may be freed :(
-    _enum_str_store = []
+    # Create an enum list from node items
+    def node_enum_items(self, context):
+        enum_items = []
+        for index, item in enumerate(nodeitems_utils.node_items_iter(context)):
+            nodetype = getattr(bpy.types, item.nodetype, None)
+            if nodetype:
+                enum_items.append((str(index), item.label, nodetype.bl_rna.description, index))
+        return enum_items
 
-    # XXX this should be called 'node_type' but the operator search
+    # Look up the item based on index
+    def find_node_item(self, context):
+        for index, item in enumerate(nodeitems_utils.node_items_iter(context)):
+            if str(index) == self.type:
+                return item
+        return None
+
+    # XXX this should be called 'node_item' but the operator search
     # property is hardcoded to 'type' by a hack in bpy_operator_wrap.c ...
     type = EnumProperty(
             name="Node Type",
             description="Node type",
-            items=node_type_items_cb,
+            items=node_enum_items,
             )
 
     def execute(self, context):
-        for index, item in enumerate(node_items_iter(context)):
-            if str(index) == self.type:
-                node = self.create_node(context, item[0].bl_rna.identifier)
-                for prop, value in item[3].items():
-                    setattr(node, prop, value)
-                break
-        return {'FINISHED'}
+        item = self.find_node_item(context)
+        if item:
+            # apply settings from the node item
+            for setting in item.settings.items():
+                ops = self.settings.add()
+                ops.name = setting[0]
+                ops.value = setting[1]
+
+            self.create_node(context, item.nodetype)
+
+            if self.use_transform:
+                bpy.ops.transform.translate('INVOKE_DEFAULT')
+
+            return {'FINISHED'}
+        else:
+            return {'CANCELLED'}
 
     def invoke(self, context, event):
         self.store_mouse_cursor(context, event)