Moved a couple of common properties into the NodeAddOperator base class to avoid repetitive code. A new operator node_add_and_link is another variant that first creates a node and them connects a specific socket to an existing one (defined by context pointer).

1 files changed, 70 insertions, 42 deletions

diff --git a/release/scripts/startup/bl_operators/node.py b/release/scripts/startup/bl_operators/node.py
index 1071f135c5f..d658350a5c3 100644
--- a/release/scripts/startup/bl_operators/node.py
+++ b/release/scripts/startup/bl_operators/node.py
@@ -20,11 +20,34 @@
 
 import bpy
 from bpy.types import Operator, PropertyGroup
-from bpy.props import BoolProperty, CollectionProperty, EnumProperty, StringProperty
+from bpy.props import BoolProperty, CollectionProperty, EnumProperty, IntProperty, StringProperty
 
 
+class NodeSetting(PropertyGroup):
+    value = StringProperty(
+            name="Value",
+            description="Python expression to be evaluated as the initial node setting",
+            default="",
+            )
+
 # Base class for node 'Add' operators
 class NodeAddOperator():
+
+    type = StringProperty(
+            name="Node Type",
+            description="Node type",
+            )
+    use_transform = BoolProperty(
+            name="Use Transform",
+            description="Start transform operator after inserting the node",
+            default=False,
+            )
+    settings = CollectionProperty(
+            name="Settings",
+            description="Settings to be applied on the newly created node",
+            type=NodeSetting,
+            )
+
     @staticmethod
     def store_mouse_cursor(context, event):
         space = context.space_data
@@ -38,7 +61,7 @@ class NodeAddOperator():
         else:
             space.cursor_location = tree.view_center
 
-    def create_node(self, context, node_type):
+    def create_node(self, context):
         space = context.space_data
         tree = space.edit_tree
 
@@ -46,7 +69,18 @@ class NodeAddOperator():
         for n in tree.nodes:
             n.select = False
 
-        node = tree.nodes.new(type=node_type)
+        node = tree.nodes.new(type=self.type)
+
+        for setting in self.settings:
+            # XXX catch exceptions here?
+            value = eval(setting.value)
+                
+            try:
+                setattr(node, setting.name, value)
+            except AttributeError as e:
+                self.report({'ERROR_INVALID_INPUT'}, "Node has no attribute "+setting.name)
+                print (str(e))
+                # Continue despite invalid attribute
 
         if space.use_hidden_preview:
             node.show_preview = False
@@ -62,18 +96,22 @@ class NodeAddOperator():
         # needs active node editor and a tree to add nodes to
         return (space.type == 'NODE_EDITOR' and space.edit_tree)
 
+    # Default execute simply adds a node
+    def execute(self, context):
+        self.create_node(context)
+        return {'FINISHED'}
+
     # Default invoke stores the mouse position to place the node correctly
+    # and optionally invokes the transform operator
     def invoke(self, context, event):
         self.store_mouse_cursor(context, event)
-        return self.execute(context)
+        result = self.execute(context)
 
+        if self.use_transform and ('FINISHED' in result):
+            bpy.ops.transform.translate('INVOKE_DEFAULT')
+
+        return result
 
-class NodeSetting(PropertyGroup):
-    value = StringProperty(
-            name="Value",
-            description="Python expression to be evaluated as the initial node setting",
-            default="",
-            )
 
 # Simple basic operator for adding a node
 class NODE_OT_add_node(NodeAddOperator, Operator):
@@ -81,45 +119,35 @@ class NODE_OT_add_node(NodeAddOperator, Operator):
     bl_idname = "node.add_node"
     bl_label = "Add Node"
 
-    type = StringProperty(
-            name="Node Type",
-            description="Node type",
-            )
-    use_transform = BoolProperty(
-            name="Use Transform",
-            description="Start transform operator after inserting the node",
-            default=False,
-            )
-    settings = CollectionProperty(
-            name="Settings",
-            description="Settings to be applied on the newly created node",
-            type=NodeSetting,
+
+# Add a node and link it to an existing socket
+class NODE_OT_add_and_link_node(NodeAddOperator, Operator):
+    '''Add a node to the active tree and link to an existing socket'''
+    bl_idname = "node.add_and_link_node"
+    bl_label = "Add and Link Node"
+
+    link_socket_index = IntProperty(
+            name="Link Socket Index",
+            description="Index of the socket to link",
             )
 
     def execute(self, context):
-        node = self.create_node(context, self.type)
-
-        for setting in self.settings:
-            # XXX catch exceptions here?
-            value = eval(setting.value)
-                
-            try:
-                setattr(node, setting.name, value)
-            except AttributeError as e:
-                self.report({'ERROR_INVALID_INPUT'}, "Node has no attribute "+setting.name)
-                print (str(e))
-                # Continue despite invalid attribute
+        space = context.space_data
+        ntree = space.edit_tree
 
-        return {'FINISHED'}
+        node = self.create_node(context)
+        if not node:
+            return {'CANCELLED'}
 
-    def invoke(self, context, event):
-        self.store_mouse_cursor(context, event)
-        result = self.execute(context)
+        to_socket = getattr(context, "link_to_socket", None)
+        if to_socket:
+            ntree.links.new(node.outputs[self.link_socket_index], to_socket)
 
-        if self.use_transform and ('FINISHED' in result):
-            bpy.ops.transform.translate('INVOKE_DEFAULT')
+        from_socket = getattr(context, "link_from_socket", None)
+        if from_socket:
+            ntree.links.new(from_socket, node.inputs[self.link_socket_index])
 
-        return result
+        return {'FINISHED'}
 
 
 def node_classes_iter(base=bpy.types.Node):