release/scripts/startup/bl_operators/node.py

   1 # ##### BEGIN GPL LICENSE BLOCK #####
   2 #
   3 #  This program is free software; you can redistribute it and/or
   4 #  modify it under the terms of the GNU General Public License
   5 #  as published by the Free Software Foundation; either version 2
   6 #  of the License, or (at your option) any later version.
   7 #
   8 #  This program is distributed in the hope that it will be useful,
   9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 #  GNU General Public License for more details.
  12 #
  13 #  You should have received a copy of the GNU General Public License
  14 #  along with this program; if not, write to the Free Software Foundation,
  15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  16 #
  17 # ##### END GPL LICENSE BLOCK #####
  18
  19 # <pep8-80 compliant>
  20
  21 import bpy
  22 from bpy.types import Operator
  23 from bpy.props import BoolProperty, EnumProperty, StringProperty
  24
  25
  26 # Base class for node 'Add' operators
  27 class NodeAddOperator():
  28     @staticmethod
  29     def store_mouse_cursor(context, event):
  30         space = context.space_data
  31         v2d = context.region.view2d
  32
  33         # convert mouse position to the View2D for later node placement
  34         space.cursor_location = v2d.region_to_view(event.mouse_region_x,
  35                                                    event.mouse_region_y)
  36
  37     def create_node(self, context, node_type):
  38         space = context.space_data
  39         tree = space.edit_tree
  40
  41         # select only the new node
  42         for n in tree.nodes:
  43             n.select = False
  44
  45         node = tree.nodes.new(type=node_type)
  46
  47         if space.use_hidden_preview:
  48             node.show_preview = False
  49
  50         node.select = True
  51         tree.nodes.active = node
  52         node.location = space.cursor_location
  53         return node
  54
  55     @classmethod
  56     def poll(cls, context):
  57         space = context.space_data
  58         # needs active node editor and a tree to add nodes to
  59         return (space.type == 'NODE_EDITOR' and space.edit_tree)
  60
  61     # Default invoke stores the mouse position to place the node correctly
  62     def invoke(self, context, event):
  63         self.store_mouse_cursor(context, event)
  64         return self.execute(context)
  65
  66
  67 # Simple basic operator for adding a node
  68 class NODE_OT_add_node(NodeAddOperator, Operator):
  69     '''Add a node to the active tree'''
  70     bl_idname = "node.add_node"
  71     bl_label = "Add Node"
  72
  73     type = StringProperty(
  74             name="Node Type",
  75             description="Node type",
  76             )
  77     # optional group tree parameter for group nodes
  78     group_tree = StringProperty(
  79             name="Group tree",
  80             description="Group node tree name",
  81             )
  82     use_transform = BoolProperty(
  83             name="Use Transform",
  84             description="Start transform operator after inserting the node",
  85             default=False,
  86             )
  87
  88     def execute(self, context):
  89         node = self.create_node(context, self.type)
  90
  91         # set the node group tree of a group node
  92         if self.properties.is_property_set('group_tree'):
  93             node.node_tree = bpy.data.node_groups[self.group_tree]
  94
  95         return {'FINISHED'}
  96
  97     def invoke(self, context, event):
  98         self.store_mouse_cursor(context, event)
  99         result = self.execute(context)
 100         if self.use_transform and ('FINISHED' in result):
 101             return bpy.ops.transform.translate('INVOKE_DEFAULT')
 102         else:
 103             return result
 104
 105
 106 def node_classes_iter(base=bpy.types.Node):
 107     """
 108     Yields all true node classes by checking for the is_registered_node_type classmethod.
 109     Node types can use specialized subtypes of bpy.types.Node, which are not usable
 110     nodes themselves (e.g. CompositorNode).
 111     """
 112     if base.is_registered_node_type():
 113         yield base
 114     for subclass in base.__subclasses__():
 115         for node_class in node_classes_iter(subclass):
 116             yield node_class
 117
 118
 119 def node_class_items_iter(node_class, context):
 120     identifier = node_class.bl_rna.identifier
 121     # XXX Checking for explicit group node types is stupid.
 122     # This should be replaced by a generic system of generating
 123     # node items via callback.
 124     # Group node_tree pointer should also use a poll function to filter the library list,
 125     # but cannot do that without a node instance here. A node callback could just use the internal poll function.
 126     if identifier in {'ShaderNodeGroup', 'CompositorNodeGroup', 'TextureNodeGroup'}:
 127         tree_idname = context.space_data.edit_tree.bl_idname
 128         for group in bpy.data.node_groups:
 129             if group.bl_idname == tree_idname:
 130                 # XXX empty string should be replaced by description from tree
 131                 yield (group.name, "", {"node_tree": group})
 132     else:
 133         yield (node_class.bl_rna.name, node_class.bl_rna.description, {})
 134
 135
 136 def node_items_iter(context):
 137     snode = context.space_data
 138     if not snode:
 139         return
 140     tree = snode.edit_tree
 141     if not tree:
 142         return
 143
 144     for node_class in node_classes_iter():
 145         if node_class.poll(tree):
 146             for item in node_class_items_iter(node_class, context):
 147                 yield (node_class,) + item
 148
 149
 150 # Create an enum list from node class items
 151 def node_type_items_cb(self, context):
 152     return [(str(index), item[1], item[2]) for index, item in enumerate(node_items_iter(context))]
 153
 154
 155 class NODE_OT_add_search(NodeAddOperator, Operator):
 156     '''Add a node to the active tree'''
 157     bl_idname = "node.add_search"
 158     bl_label = "Search and Add Node"
 159     bl_options = {'REGISTER', 'UNDO'}
 160
 161     # XXX this should be called 'node_type' but the operator search
 162     # property is hardcoded to 'type' by a hack in bpy_operator_wrap.c ...
 163     type = EnumProperty(
 164             name="Node Type",
 165             description="Node type",
 166             items=node_type_items_cb,
 167             )
 168
 169     def execute(self, context):
 170         for index, item in enumerate(node_items_iter(context)):
 171             if str(index) == self.type:
 172                 node = self.create_node(context, item[0].bl_rna.identifier)
 173                 for prop, value in item[3].items():
 174                     setattr(node, prop, value)
 175                 break
 176         return {'FINISHED'}
 177
 178     def invoke(self, context, event):
 179         self.store_mouse_cursor(context, event)
 180         # Delayed execution in the search popup
 181         context.window_manager.invoke_search_popup(self)
 182         return {'CANCELLED'}
 183
 184
 185 # Simple basic operator for adding a node without further initialization
 186 class NODE_OT_add_node(NodeAddOperator, bpy.types.Operator):
 187     '''Add a node to the active tree'''
 188     bl_idname = "node.add_node"
 189     bl_label = "Add Node"
 190
 191     type = StringProperty(name="Node Type", description="Node type")
 192
 193     def execute(self, context):
 194         node = self.create_node(context, self.type)
 195         return {'FINISHED'}
 196
 197
 198 class NODE_OT_add_group_node(NodeAddOperator, bpy.types.Operator):
 199     '''Add a group node to the active tree'''
 200     bl_idname = "node.add_group_node"
 201     bl_label = "Add Group Node"
 202
 203     type = StringProperty(name="Node Type", description="Node type")
 204     grouptree = StringProperty(name="Group tree", description="Group node tree name")
 205
 206     def execute(self, context):
 207         node = self.create_node(context, self.type)
 208         node.node_tree = bpy.data.node_groups[self.grouptree]
 209
 210         return {'FINISHED'}
 211
 212
 213 class NODE_OT_collapse_hide_unused_toggle(Operator):
 214     '''Toggle collapsed nodes and hide unused sockets'''
 215     bl_idname = "node.collapse_hide_unused_toggle"
 216     bl_label = "Collapse and Hide Unused Sockets"
 217     bl_options = {'REGISTER', 'UNDO'}
 218
 219     @classmethod
 220     def poll(cls, context):
 221         space = context.space_data
 222         # needs active node editor and a tree
 223         return (space.type == 'NODE_EDITOR' and space.edit_tree)
 224
 225     def execute(self, context):
 226         space = context.space_data
 227         tree = space.edit_tree
 228
 229         for node in tree.nodes:
 230             if node.select:
 231                 hide = (not node.hide)
 232
 233                 node.hide = hide
 234                 # Note: connected sockets are ignored internally
 235                 for socket in node.inputs:
 236                     socket.hide = hide
 237                 for socket in node.outputs:
 238                     socket.hide = hide
 239
 240         return {'FINISHED'}
 241
 242
 243 class NODE_OT_tree_path_parent(Operator):
 244     '''Go to parent node tree'''
 245     bl_idname = "node.tree_path_parent"
 246     bl_label = "Parent Node Tree"
 247     bl_options = {'REGISTER', 'UNDO'}
 248
 249     @classmethod
 250     def poll(cls, context):
 251         space = context.space_data
 252         # needs active node editor and a tree
 253         return (space.type == 'NODE_EDITOR' and len(space.path) > 1)
 254
 255     def execute(self, context):
 256         space = context.space_data
 257
 258         space.path.pop()
 259
 260         return {'FINISHED'}