5c7f3c3f4b27157919bd63e556d04809a1a99888
[blender-staging.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 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                 yield (group.name, "", {"node_tree":group}) # XXX empty string should be replaced by description from tree
131     else:
132         yield (node_class.bl_rna.name, node_class.bl_rna.description, {})
133
134
135 def node_items_iter(context):
136     snode = context.space_data
137     if not snode:
138         return
139     tree = snode.edit_tree
140     if not tree:
141         return
142
143     for node_class in node_classes_iter():
144         if node_class.poll(tree):
145             for item in node_class_items_iter(node_class, context):
146                 yield (node_class,) + item
147
148
149 # Create an enum list from node class items
150 def node_type_items_cb(self, context):
151     return [(str(index), item[1], item[2]) for index, item in enumerate(node_items_iter(context))]
152
153
154 class NODE_OT_add_search(NodeAddOperator, Operator):
155     '''Add a node to the active tree'''
156     bl_idname = "node.add_search"
157     bl_label = "Search and Add Node"
158     bl_options = {'REGISTER', 'UNDO'}
159
160     # XXX this should be called 'node_type' but the operator search
161     # property is hardcoded to 'type' by a hack in bpy_operator_wrap.c ...
162     type = EnumProperty(
163             name="Node Type",
164             description="Node type",
165             items=node_type_items_cb,
166             )
167
168     def execute(self, context):
169         for index, item in enumerate(node_items_iter(context)):
170             if str(index) == self.type:
171                 node = self.create_node(context, item[0].bl_rna.identifier)
172                 for prop,value in item[3].items():
173                     setattr(node, prop, value)
174                 break
175         return {'FINISHED'}
176
177     def invoke(self, context, event):
178         self.store_mouse_cursor(context, event)
179         # Delayed execution in the search popup
180         context.window_manager.invoke_search_popup(self)
181         return {'CANCELLED'}
182
183
184 # Simple basic operator for adding a node without further initialization
185 class NODE_OT_add_node(NodeAddOperator, bpy.types.Operator):
186     '''Add a node to the active tree'''
187     bl_idname = "node.add_node"
188     bl_label = "Add Node"
189
190     type = StringProperty(name="Node Type", description="Node type")
191
192     def execute(self, context):
193         node = self.create_node(context, self.type)
194         return {'FINISHED'}
195
196
197 class NODE_OT_add_group_node(NodeAddOperator, bpy.types.Operator):
198     '''Add a group node to the active tree'''
199     bl_idname = "node.add_group_node"
200     bl_label = "Add Group Node"
201
202     type = StringProperty(name="Node Type", description="Node type")
203     grouptree = StringProperty(name="Group tree", description="Group node tree name")
204
205     def execute(self, context):
206         node = self.create_node(context, self.type)
207         node.node_tree = bpy.data.node_groups[self.grouptree]
208
209         return {'FINISHED'}
210
211
212 class NODE_OT_collapse_hide_unused_toggle(Operator):
213     '''Toggle collapsed nodes and hide unused sockets'''
214     bl_idname = "node.collapse_hide_unused_toggle"
215     bl_label = "Collapse and Hide Unused Sockets"
216     bl_options = {'REGISTER', 'UNDO'}
217
218     @classmethod
219     def poll(cls, context):
220         space = context.space_data
221         # needs active node editor and a tree
222         return (space.type == 'NODE_EDITOR' and space.edit_tree)
223
224     def execute(self, context):
225         space = context.space_data
226         tree = space.edit_tree
227
228         for node in tree.nodes:
229             if node.select:
230                 hide = (not node.hide)
231
232                 node.hide = hide
233                 # Note: connected sockets are ignored internally
234                 for socket in node.inputs:
235                     socket.hide = hide
236                 for socket in node.outputs:
237                     socket.hide = hide
238
239         return {'FINISHED'}
240
241
242 class NODE_OT_tree_path_parent(Operator):
243     '''Go to parent node tree'''
244     bl_idname = "node.tree_path_parent"
245     bl_label = "Parent Node Tree"
246     bl_options = {'REGISTER', 'UNDO'}
247
248     @classmethod
249     def poll(cls, context):
250         space = context.space_data
251         # needs active node editor and a tree
252         return (space.type == 'NODE_EDITOR' and len(space.path) > 1)
253
254     def execute(self, context):
255         space = context.space_data
256
257         space.path.pop()
258
259         return {'FINISHED'}
260