Fix #35640, part 2. Check id.lib in poll functions for operators which do critical...
[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, nodeitems_utils
22 from bpy.types import Operator, PropertyGroup
23 from bpy.props import BoolProperty, CollectionProperty, EnumProperty, IntProperty, StringProperty
24
25
26 class NodeSetting(PropertyGroup):
27     value = StringProperty(
28             name="Value",
29             description="Python expression to be evaluated as the initial node setting",
30             default="",
31             )
32
33 # Base class for node 'Add' operators
34 class NodeAddOperator():
35
36     type = StringProperty(
37             name="Node Type",
38             description="Node type",
39             )
40     use_transform = BoolProperty(
41             name="Use Transform",
42             description="Start transform operator after inserting the node",
43             default=False,
44             )
45     settings = CollectionProperty(
46             name="Settings",
47             description="Settings to be applied on the newly created node",
48             type=NodeSetting,
49             options={'SKIP_SAVE'},
50             )
51
52     @staticmethod
53     def store_mouse_cursor(context, event):
54         space = context.space_data
55         v2d = context.region.view2d
56         tree = space.edit_tree
57
58         # convert mouse position to the View2D for later node placement
59         if context.region.type == 'WINDOW':
60             space.cursor_location = v2d.region_to_view(event.mouse_region_x,
61                                                    event.mouse_region_y)
62         else:
63             space.cursor_location = tree.view_center
64
65     # XXX explicit node_type argument is usually not necessary, but required to make search operator work:
66     # add_search has to override the 'type' property since it's hardcoded in bpy_operator_wrap.c ...
67     def create_node(self, context, node_type=None):
68         space = context.space_data
69         tree = space.edit_tree
70
71         if node_type is None:
72             node_type = self.type
73
74         # select only the new node
75         for n in tree.nodes:
76             n.select = False
77
78         node = tree.nodes.new(type=node_type)
79
80         for setting in self.settings:
81             # XXX catch exceptions here?
82             value = eval(setting.value)
83                 
84             try:
85                 setattr(node, setting.name, value)
86             except AttributeError as e:
87                 self.report({'ERROR_INVALID_INPUT'}, "Node has no attribute "+setting.name)
88                 print (str(e))
89                 # Continue despite invalid attribute
90
91         if space.use_hidden_preview:
92             node.show_preview = False
93
94         node.select = True
95         tree.nodes.active = node
96         node.location = space.cursor_location
97         return node
98
99     @classmethod
100     def poll(cls, context):
101         space = context.space_data
102         # needs active node editor and a tree to add nodes to
103         return (space.type == 'NODE_EDITOR' and space.edit_tree and not space.edit_tree.library)
104
105     # Default execute simply adds a node
106     def execute(self, context):
107         self.create_node(context)
108         return {'FINISHED'}
109
110     # Default invoke stores the mouse position to place the node correctly
111     # and optionally invokes the transform operator
112     def invoke(self, context, event):
113         self.store_mouse_cursor(context, event)
114         result = self.execute(context)
115
116         if self.use_transform and ('FINISHED' in result):
117             bpy.ops.transform.translate('INVOKE_DEFAULT')
118
119         return result
120
121
122 # Simple basic operator for adding a node
123 class NODE_OT_add_node(NodeAddOperator, Operator):
124     '''Add a node to the active tree'''
125     bl_idname = "node.add_node"
126     bl_label = "Add Node"
127     bl_options = {'REGISTER', 'UNDO'}
128
129
130 # Add a node and link it to an existing socket
131 class NODE_OT_add_and_link_node(NodeAddOperator, Operator):
132     '''Add a node to the active tree and link to an existing socket'''
133     bl_idname = "node.add_and_link_node"
134     bl_label = "Add and Link Node"
135     bl_options = {'REGISTER', 'UNDO'}
136
137     link_socket_index = IntProperty(
138             name="Link Socket Index",
139             description="Index of the socket to link",
140             )
141
142     def execute(self, context):
143         space = context.space_data
144         ntree = space.edit_tree
145
146         node = self.create_node(context)
147         if not node:
148             return {'CANCELLED'}
149
150         to_socket = getattr(context, "link_to_socket", None)
151         if to_socket:
152             ntree.links.new(node.outputs[self.link_socket_index], to_socket)
153
154         from_socket = getattr(context, "link_from_socket", None)
155         if from_socket:
156             ntree.links.new(from_socket, node.inputs[self.link_socket_index])
157
158         return {'FINISHED'}
159
160
161 class NODE_OT_add_search(NodeAddOperator, Operator):
162     '''Add a node to the active tree'''
163     bl_idname = "node.add_search"
164     bl_label = "Search and Add Node"
165     bl_options = {'REGISTER', 'UNDO'}
166     bl_property = "node_item"
167
168     _enum_item_hack = []
169
170     # Create an enum list from node items
171     def node_enum_items(self, context):
172         enum_items = NODE_OT_add_search._enum_item_hack
173         enum_items.clear()
174
175         for index, item in enumerate(nodeitems_utils.node_items_iter(context)):
176             nodetype = getattr(bpy.types, item.nodetype, None)
177             if nodetype:
178                 enum_items.append((str(index), item.label, nodetype.bl_rna.description, index))
179         return enum_items
180
181     # Look up the item based on index
182     def find_node_item(self, context):
183         node_item = int(self.node_item)
184         for index, item in enumerate(nodeitems_utils.node_items_iter(context)):
185             if index == node_item:
186                 return item
187         return None
188
189     node_item = EnumProperty(
190             name="Node Type",
191             description="Node type",
192             items=node_enum_items,
193             )
194
195     def execute(self, context):
196         item = self.find_node_item(context)
197
198         # no need to keep
199         self._enum_item_hack.clear()
200         
201         if item:
202             # apply settings from the node item
203             for setting in item.settings.items():
204                 ops = self.settings.add()
205                 ops.name = setting[0]
206                 ops.value = setting[1]
207
208             self.create_node(context, item.nodetype)
209
210             if self.use_transform:
211                 bpy.ops.transform.translate('INVOKE_DEFAULT')
212
213             return {'FINISHED'}
214         else:
215             return {'CANCELLED'}
216
217     def invoke(self, context, event):
218         self.store_mouse_cursor(context, event)
219         # Delayed execution in the search popup
220         context.window_manager.invoke_search_popup(self)
221         return {'CANCELLED'}
222
223
224 class NODE_OT_collapse_hide_unused_toggle(Operator):
225     '''Toggle collapsed nodes and hide unused sockets'''
226     bl_idname = "node.collapse_hide_unused_toggle"
227     bl_label = "Collapse and Hide Unused Sockets"
228     bl_options = {'REGISTER', 'UNDO'}
229
230     @classmethod
231     def poll(cls, context):
232         space = context.space_data
233         # needs active node editor and a tree
234         return (space.type == 'NODE_EDITOR' and space.edit_tree and not space.edit_tree.library)
235
236     def execute(self, context):
237         space = context.space_data
238         tree = space.edit_tree
239
240         for node in tree.nodes:
241             if node.select:
242                 hide = (not node.hide)
243
244                 node.hide = hide
245                 # Note: connected sockets are ignored internally
246                 for socket in node.inputs:
247                     socket.hide = hide
248                 for socket in node.outputs:
249                     socket.hide = hide
250
251         return {'FINISHED'}
252
253
254 class NODE_OT_tree_path_parent(Operator):
255     '''Go to parent node tree'''
256     bl_idname = "node.tree_path_parent"
257     bl_label = "Parent Node Tree"
258     bl_options = {'REGISTER', 'UNDO'}
259
260     @classmethod
261     def poll(cls, context):
262         space = context.space_data
263         # needs active node editor and a tree
264         return (space.type == 'NODE_EDITOR' and len(space.path) > 1)
265
266     def execute(self, context):
267         space = context.space_data
268
269         space.path.pop()
270
271         return {'FINISHED'}