svn merge ^/trunk/blender -r49945:49947
[blender.git] / 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 EnumProperty
24
25 # XXX These node item lists should actually be generated by a callback at operator execution time (see node_type_items below),
26 # using the active node tree from the context. Due to a difficult bug in bpy this is not possible (item list memory gets freed too early),
27 # so for now just copy the static item lists to these global variables.
28 #
29 # In the custom_nodes branch, the static per-tree-type node items are replaced by a single independent type list anyway (with a poll function
30 # to limit node types to the respective trees). So this workaround is only temporary.
31
32 # lazy init
33 node_type_items_dict = {}
34
35 # Prefixes used to distinguish base node types and node groups
36 node_type_prefix = 'NODE_'
37 node_group_prefix = 'GROUP_'
38
39
40 # Generate a list of enum items for a given node class
41 # Copy existing type enum, adding a prefix to distinguish from node groups
42 # Skip the base node group type, node groups will be added below for all existing group trees
43 def node_type_items(node_class):
44     return [(node_type_prefix + item.identifier, item.name, item.description)
45                     for item in node_class.bl_rna.properties['type'].enum_items if item.identifier != 'GROUP']
46
47
48 # Generate items for node group types
49 # Filter by the given tree_type
50 # Node group trees don't have a description property yet (could add this as a custom property though)
51 def node_group_items(tree_type):
52     return [(node_group_prefix + group.name, group.name, '')
53                     for group in bpy.data.node_groups if group.type == tree_type]
54
55
56 # Returns the enum item list for the edited tree in the context
57 def node_type_items_cb(self, context):
58     snode = context.space_data
59     if not snode:
60         return ()
61     tree = snode.edit_tree
62     if not tree:
63         return ()
64
65     # Lists of basic node types for each
66     if not node_type_items_dict:
67         node_type_items_dict.update({
68             'SHADER': node_type_items(bpy.types.ShaderNode),
69             'COMPOSITING': node_type_items(bpy.types.CompositorNode),
70             'TEXTURE': node_type_items(bpy.types.TextureNode),
71             })
72
73     # XXX Does not work correctly, see comment above
74     #return [(item.identifier, item.name, item.description, item.value) for item in tree.nodes.bl_rna.functions['new'].parameters['type'].enum_items]
75
76     if tree.type in node_type_items_dict:
77         return node_type_items_dict[tree.type] + node_group_items(tree.type)
78     else:
79         return ()
80
81
82 class NODE_OT_add_search(Operator):
83     '''Add a node to the active tree'''
84     bl_idname = "node.add_search"
85     bl_label = "Search and Add Node"
86     bl_options = {'REGISTER', 'UNDO'}
87
88     # XXX this should be called 'node_type' but the operator search property is hardcoded to 'type' by a hack in bpy_operator_wrap.c ...
89     type = EnumProperty(
90             name="Node Type",
91             description="Node type",
92             items=node_type_items_cb,
93             )
94
95     _node_type_items_dict = None
96
97     def create_node(self, context):
98         space = context.space_data
99         tree = space.edit_tree
100
101         # Enum item identifier has an additional prefix to distinguish base node types from node groups
102         item = self.type
103         if (item.startswith(node_type_prefix)):
104             # item means base node type
105             node = tree.nodes.new(type=item[len(node_type_prefix):])
106         elif (item.startswith(node_group_prefix)):
107             # item means node group type
108             node = tree.nodes.new(type='GROUP', group=bpy.data.node_groups[item[len(node_group_prefix):]])
109         else:
110             return None
111
112         for n in tree.nodes:
113             if n == node:
114                 node.select = True
115                 tree.nodes.active = node
116             else:
117                 node.select = False
118         node.location = space.cursor_location
119         return node
120
121     @classmethod
122     def poll(cls, context):
123         space = context.space_data
124         # needs active node editor and a tree to add nodes to
125         return (space.type == 'NODE_EDITOR' and space.edit_tree)
126
127     def execute(self, context):
128         self.create_node(context)
129         return {'FINISHED'}
130
131     def invoke(self, context, event):
132         space = context.space_data
133         v2d = context.region.view2d
134
135         # convert mouse position to the View2D for later node placement
136         space.cursor_location = v2d.region_to_view(event.mouse_region_x, event.mouse_region_y)
137
138         context.window_manager.invoke_search_popup(self)
139         return {'CANCELLED'}
140
141
142 class NODE_OT_collapse_hide_unused_toggle(Operator):
143     '''Toggle collapsed nodes and hide unused sockets'''
144     bl_idname = "node.collapse_hide_unused_toggle"
145     bl_label = "Collapse and Hide Unused Sockets"
146     bl_options = {'REGISTER', 'UNDO'}
147
148     @classmethod
149     def poll(cls, context):
150         space = context.space_data
151         # needs active node editor and a tree
152         return (space.type == 'NODE_EDITOR' and space.edit_tree)
153
154     def execute(self, context):
155         space = context.space_data
156         tree = space.edit_tree
157
158         for node in tree.nodes:
159             if node.select:
160                 hide = (not node.hide)
161
162                 node.hide = hide
163                 # Note: connected sockets are ignored internally
164                 for socket in node.inputs:
165                     socket.hide = hide
166                 for socket in node.outputs:
167                     socket.hide = hide
168
169         return {'FINISHED'}