Cleanup: pep8 & redundant vars
[blender.git] / release / scripts / startup / nodeitems_builtins.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 compliant>
20 import bpy
21 import nodeitems_utils
22 from nodeitems_utils import NodeCategory, NodeItem, NodeItemCustom
23
24
25 # Subclasses for standard node types
26
27 class CompositorNodeCategory(NodeCategory):
28     @classmethod
29     def poll(cls, context):
30         return (context.space_data.tree_type == 'CompositorNodeTree')
31
32
33 class ShaderNewNodeCategory(NodeCategory):
34     @classmethod
35     def poll(cls, context):
36         return (context.space_data.tree_type == 'ShaderNodeTree' and
37                 context.scene.render.use_shading_nodes)
38
39
40 class ShaderOldNodeCategory(NodeCategory):
41     @classmethod
42     def poll(cls, context):
43         return (context.space_data.tree_type == 'ShaderNodeTree' and
44                 not context.scene.render.use_shading_nodes)
45
46
47 class TextureNodeCategory(NodeCategory):
48     @classmethod
49     def poll(cls, context):
50         return context.space_data.tree_type == 'TextureNodeTree'
51
52
53 # menu entry for node group tools
54 def group_tools_draw(self, layout, context):
55     layout.operator("node.group_make")
56     layout.operator("node.group_ungroup")
57     layout.separator()
58
59 # maps node tree type to group node type
60 node_tree_group_type = {
61     'CompositorNodeTree': 'CompositorNodeGroup',
62     'ShaderNodeTree': 'ShaderNodeGroup',
63     'TextureNodeTree': 'TextureNodeGroup',
64     }
65
66
67 # generic node group items generator for shader, compositor and texture node groups
68 def node_group_items(context):
69     space = context.space_data
70     if not space:
71         return
72     ntree = space.edit_tree
73     if not ntree:
74         return
75
76     yield NodeItemCustom(draw=group_tools_draw)
77
78     def contains_group(nodetree, group):
79         if nodetree == group:
80             return True
81         else:
82             for node in nodetree.nodes:
83                 if node.bl_idname in node_tree_group_type.values() and node.node_tree is not None:
84                     if contains_group(node.node_tree, group):
85                         return True
86         return False
87
88     for group in context.blend_data.node_groups:
89         if group.bl_idname != ntree.bl_idname:
90             continue
91         # filter out recursive groups
92         if contains_group(group, ntree):
93             continue
94
95         yield NodeItem(node_tree_group_type[group.bl_idname],
96                        group.name,
97                        {"node_tree": "bpy.data.node_groups[%r]" % group.name})
98
99
100 # only show input/output nodes inside node groups
101 def group_input_output_item_poll(context):
102     space = context.space_data
103     if space.edit_tree in bpy.data.node_groups.values():
104         return True
105     return False
106
107
108 # All standard node categories currently used in nodes.
109
110 shader_node_categories = [
111     # Shader Nodes
112     ShaderOldNodeCategory("SH_INPUT", "Input", items=[
113         NodeItem("ShaderNodeMaterial"),
114         NodeItem("ShaderNodeCameraData"),
115         NodeItem("ShaderNodeLampData"),
116         NodeItem("ShaderNodeValue"),
117         NodeItem("ShaderNodeRGB"),
118         NodeItem("ShaderNodeTexture"),
119         NodeItem("ShaderNodeGeometry"),
120         NodeItem("ShaderNodeExtendedMaterial"),
121         NodeItem("NodeGroupInput", poll=group_input_output_item_poll),
122         ]),
123     ShaderOldNodeCategory("SH_OUTPUT", "Output", items=[
124         NodeItem("ShaderNodeOutput"),
125         NodeItem("NodeGroupOutput", poll=group_input_output_item_poll),
126         ]),
127     ShaderOldNodeCategory("SH_OP_COLOR", "Color", items=[
128         NodeItem("ShaderNodeMixRGB"),
129         NodeItem("ShaderNodeRGBCurve"),
130         NodeItem("ShaderNodeInvert"),
131         NodeItem("ShaderNodeHueSaturation"),
132         ]),
133     ShaderOldNodeCategory("SH_OP_VECTOR", "Vector", items=[
134         NodeItem("ShaderNodeNormal"),
135         NodeItem("ShaderNodeMapping"),
136         NodeItem("ShaderNodeVectorCurve"),
137         ]),
138     ShaderOldNodeCategory("SH_CONVERTOR", "Converter", items=[
139         NodeItem("ShaderNodeValToRGB"),
140         NodeItem("ShaderNodeRGBToBW"),
141         NodeItem("ShaderNodeMath"),
142         NodeItem("ShaderNodeVectorMath"),
143         NodeItem("ShaderNodeSqueeze"),
144         NodeItem("ShaderNodeSeparateRGB"),
145         NodeItem("ShaderNodeCombineRGB"),
146         NodeItem("ShaderNodeSeparateHSV"),
147         NodeItem("ShaderNodeCombineHSV"),
148         ]),
149     ShaderOldNodeCategory("SH_GROUP", "Group", items=node_group_items),
150     ShaderOldNodeCategory("SH_LAYOUT", "Layout", items=[
151         NodeItem("NodeFrame"),
152         NodeItem("NodeReroute"),
153         ]),
154
155     # New Shader Nodes (Cycles)
156     ShaderNewNodeCategory("SH_NEW_INPUT", "Input", items=[
157         NodeItem("ShaderNodeTexCoord"),
158         NodeItem("ShaderNodeAttribute"),
159         NodeItem("ShaderNodeLightPath"),
160         NodeItem("ShaderNodeFresnel"),
161         NodeItem("ShaderNodeLayerWeight"),
162         NodeItem("ShaderNodeRGB"),
163         NodeItem("ShaderNodeValue"),
164         NodeItem("ShaderNodeTangent"),
165         NodeItem("ShaderNodeNewGeometry"),
166         NodeItem("ShaderNodeWireframe"),
167         NodeItem("ShaderNodeObjectInfo"),
168         NodeItem("ShaderNodeHairInfo"),
169         NodeItem("ShaderNodeParticleInfo"),
170         NodeItem("ShaderNodeCameraData"),
171         NodeItem("ShaderNodeUVMap"),
172         NodeItem("NodeGroupInput", poll=group_input_output_item_poll),
173         ]),
174     ShaderNewNodeCategory("SH_NEW_OUTPUT", "Output", items=[
175         NodeItem("ShaderNodeOutputMaterial"),
176         NodeItem("ShaderNodeOutputLamp"),
177         NodeItem("ShaderNodeOutputWorld"),
178         NodeItem("NodeGroupOutput", poll=group_input_output_item_poll),
179         ]),
180     ShaderNewNodeCategory("SH_NEW_SHADER", "Shader", items=[
181         NodeItem("ShaderNodeMixShader"),
182         NodeItem("ShaderNodeAddShader"),
183         NodeItem("ShaderNodeBsdfDiffuse"),
184         NodeItem("ShaderNodeBsdfGlossy"),
185         NodeItem("ShaderNodeBsdfTransparent"),
186         NodeItem("ShaderNodeBsdfRefraction"),
187         NodeItem("ShaderNodeBsdfGlass"),
188         NodeItem("ShaderNodeBsdfTranslucent"),
189         NodeItem("ShaderNodeBsdfAnisotropic"),
190         NodeItem("ShaderNodeBsdfVelvet"),
191         NodeItem("ShaderNodeBsdfToon"),
192         NodeItem("ShaderNodeSubsurfaceScattering"),
193         NodeItem("ShaderNodeEmission"),
194         NodeItem("ShaderNodeBsdfHair"),
195         NodeItem("ShaderNodeBackground"),
196         NodeItem("ShaderNodeAmbientOcclusion"),
197         NodeItem("ShaderNodeHoldout"),
198         NodeItem("ShaderNodeVolumeAbsorption"),
199         NodeItem("ShaderNodeVolumeScatter"),
200         ]),
201     ShaderNewNodeCategory("SH_NEW_TEXTURE", "Texture", items=[
202         NodeItem("ShaderNodeTexImage"),
203         NodeItem("ShaderNodeTexEnvironment"),
204         NodeItem("ShaderNodeTexSky"),
205         NodeItem("ShaderNodeTexNoise"),
206         NodeItem("ShaderNodeTexWave"),
207         NodeItem("ShaderNodeTexVoronoi"),
208         NodeItem("ShaderNodeTexMusgrave"),
209         NodeItem("ShaderNodeTexGradient"),
210         NodeItem("ShaderNodeTexMagic"),
211         NodeItem("ShaderNodeTexChecker"),
212         NodeItem("ShaderNodeTexBrick"),
213         ]),
214     ShaderNewNodeCategory("SH_NEW_OP_COLOR", "Color", items=[
215         NodeItem("ShaderNodeMixRGB"),
216         NodeItem("ShaderNodeRGBCurve"),
217         NodeItem("ShaderNodeInvert"),
218         NodeItem("ShaderNodeLightFalloff"),
219         NodeItem("ShaderNodeHueSaturation"),
220         NodeItem("ShaderNodeGamma"),
221         NodeItem("ShaderNodeBrightContrast"),
222         ]),
223     ShaderNewNodeCategory("SH_NEW_OP_VECTOR", "Vector", items=[
224         NodeItem("ShaderNodeMapping"),
225         NodeItem("ShaderNodeBump"),
226         NodeItem("ShaderNodeNormalMap"),
227         NodeItem("ShaderNodeNormal"),
228         NodeItem("ShaderNodeVectorCurve"),
229         NodeItem("ShaderNodeVectorTransform"),
230         ]),
231     ShaderNewNodeCategory("SH_NEW_CONVERTOR", "Converter", items=[
232         NodeItem("ShaderNodeMath"),
233         NodeItem("ShaderNodeValToRGB"),
234         NodeItem("ShaderNodeRGBToBW"),
235         NodeItem("ShaderNodeVectorMath"),
236         NodeItem("ShaderNodeSeparateRGB"),
237         NodeItem("ShaderNodeCombineRGB"),
238         NodeItem("ShaderNodeSeparateXYZ"),
239         NodeItem("ShaderNodeCombineXYZ"),
240         NodeItem("ShaderNodeSeparateHSV"),
241         NodeItem("ShaderNodeCombineHSV"),
242         NodeItem("ShaderNodeWavelength"),
243         NodeItem("ShaderNodeBlackbody"),
244         ]),
245     ShaderNewNodeCategory("SH_NEW_SCRIPT", "Script", items=[
246         NodeItem("ShaderNodeScript"),
247         ]),
248     ShaderNewNodeCategory("SH_NEW_GROUP", "Group", items=node_group_items),
249     ShaderNewNodeCategory("SH_NEW_LAYOUT", "Layout", items=[
250         NodeItem("NodeFrame"),
251         NodeItem("NodeReroute"),
252         ]),
253     ]
254
255 compositor_node_categories = [
256     # Compositor Nodes
257     CompositorNodeCategory("CMP_INPUT", "Input", items=[
258         NodeItem("CompositorNodeRLayers"),
259         NodeItem("CompositorNodeImage"),
260         NodeItem("CompositorNodeMovieClip"),
261         NodeItem("CompositorNodeMask"),
262         NodeItem("CompositorNodeRGB"),
263         NodeItem("CompositorNodeValue"),
264         NodeItem("CompositorNodeTexture"),
265         NodeItem("CompositorNodeBokehImage"),
266         NodeItem("CompositorNodeTime"),
267         NodeItem("CompositorNodeTrackPos"),
268         NodeItem("NodeGroupInput", poll=group_input_output_item_poll),
269         ]),
270     CompositorNodeCategory("CMP_OUTPUT", "Output", items=[
271         NodeItem("CompositorNodeComposite"),
272         NodeItem("CompositorNodeViewer"),
273         NodeItem("CompositorNodeSplitViewer"),
274         NodeItem("CompositorNodeOutputFile"),
275         NodeItem("CompositorNodeLevels"),
276         NodeItem("NodeGroupOutput", poll=group_input_output_item_poll),
277         ]),
278     CompositorNodeCategory("CMP_OP_COLOR", "Color", items=[
279         NodeItem("CompositorNodeMixRGB"),
280         NodeItem("CompositorNodeAlphaOver"),
281         NodeItem("CompositorNodeInvert"),
282         NodeItem("CompositorNodeCurveRGB"),
283         NodeItem("CompositorNodeHueSat"),
284         NodeItem("CompositorNodeColorBalance"),
285         NodeItem("CompositorNodeHueCorrect"),
286         NodeItem("CompositorNodeBrightContrast"),
287         NodeItem("CompositorNodeGamma"),
288         NodeItem("CompositorNodeColorCorrection"),
289         NodeItem("CompositorNodeTonemap"),
290         NodeItem("CompositorNodeZcombine"),
291         ]),
292     CompositorNodeCategory("CMP_CONVERTOR", "Converter", items=[
293         NodeItem("CompositorNodeMath"),
294         NodeItem("CompositorNodeValToRGB"),
295         NodeItem("CompositorNodeSetAlpha"),
296         NodeItem("CompositorNodePremulKey"),
297         NodeItem("CompositorNodeIDMask"),
298         NodeItem("CompositorNodeRGBToBW"),
299         NodeItem("CompositorNodeSepRGBA"),
300         NodeItem("CompositorNodeCombRGBA"),
301         NodeItem("CompositorNodeSepHSVA"),
302         NodeItem("CompositorNodeCombHSVA"),
303         NodeItem("CompositorNodeSepYUVA"),
304         NodeItem("CompositorNodeCombYUVA"),
305         NodeItem("CompositorNodeSepYCCA"),
306         NodeItem("CompositorNodeCombYCCA"),
307         ]),
308     CompositorNodeCategory("CMP_OP_FILTER", "Filter", items=[
309         NodeItem("CompositorNodeBlur"),
310         NodeItem("CompositorNodeBilateralblur"),
311         NodeItem("CompositorNodeDilateErode"),
312         NodeItem("CompositorNodeDespeckle"),
313         NodeItem("CompositorNodeFilter"),
314         NodeItem("CompositorNodeBokehBlur"),
315         NodeItem("CompositorNodeVecBlur"),
316         NodeItem("CompositorNodeDefocus"),
317         NodeItem("CompositorNodeGlare"),
318         NodeItem("CompositorNodeInpaint"),
319         NodeItem("CompositorNodeDBlur"),
320         NodeItem("CompositorNodePixelate"),
321         ]),
322     CompositorNodeCategory("CMP_OP_VECTOR", "Vector", items=[
323         NodeItem("CompositorNodeNormal"),
324         NodeItem("CompositorNodeMapValue"),
325         NodeItem("CompositorNodeMapRange"),
326         NodeItem("CompositorNodeNormalize"),
327         NodeItem("CompositorNodeCurveVec"),
328         ]),
329     CompositorNodeCategory("CMP_MATTE", "Matte", items=[
330         NodeItem("CompositorNodeKeying"),
331         NodeItem("CompositorNodeKeyingScreen"),
332         NodeItem("CompositorNodeChannelMatte"),
333         NodeItem("CompositorNodeColorSpill"),
334         NodeItem("CompositorNodeBoxMask"),
335         NodeItem("CompositorNodeEllipseMask"),
336         NodeItem("CompositorNodeLumaMatte"),
337         NodeItem("CompositorNodeDiffMatte"),
338         NodeItem("CompositorNodeDistanceMatte"),
339         NodeItem("CompositorNodeChromaMatte"),
340         NodeItem("CompositorNodeColorMatte"),
341         NodeItem("CompositorNodeDoubleEdgeMask"),
342         ]),
343     CompositorNodeCategory("CMP_DISTORT", "Distort", items=[
344         NodeItem("CompositorNodeScale"),
345         NodeItem("CompositorNodeLensdist"),
346         NodeItem("CompositorNodeMovieDistortion"),
347         NodeItem("CompositorNodeTranslate"),
348         NodeItem("CompositorNodeRotate"),
349         NodeItem("CompositorNodeFlip"),
350         NodeItem("CompositorNodeCrop"),
351         NodeItem("CompositorNodeDisplace"),
352         NodeItem("CompositorNodeMapUV"),
353         NodeItem("CompositorNodeTransform"),
354         NodeItem("CompositorNodeStabilize"),
355         NodeItem("CompositorNodePlaneTrackDeform"),
356         NodeItem("CompositorNodeCornerPin"),
357         ]),
358     CompositorNodeCategory("CMP_GROUP", "Group", items=node_group_items),
359     CompositorNodeCategory("CMP_LAYOUT", "Layout", items=[
360         NodeItem("NodeFrame"),
361         NodeItem("NodeReroute"),
362         NodeItem("CompositorNodeSwitch"),
363         ]),
364     ]
365
366 texture_node_categories = [
367     # Texture Nodes
368     TextureNodeCategory("TEX_INPUT", "Input", items=[
369         NodeItem("TextureNodeCurveTime"),
370         NodeItem("TextureNodeCoordinates"),
371         NodeItem("TextureNodeTexture"),
372         NodeItem("TextureNodeImage"),
373         NodeItem("NodeGroupInput", poll=group_input_output_item_poll),
374         ]),
375     TextureNodeCategory("TEX_OUTPUT", "Output", items=[
376         NodeItem("TextureNodeOutput"),
377         NodeItem("TextureNodeViewer"),
378         NodeItem("NodeGroupOutput", poll=group_input_output_item_poll),
379         ]),
380     TextureNodeCategory("TEX_OP_COLOR", "Color", items=[
381         NodeItem("TextureNodeMixRGB"),
382         NodeItem("TextureNodeCurveRGB"),
383         NodeItem("TextureNodeInvert"),
384         NodeItem("TextureNodeHueSaturation"),
385         NodeItem("TextureNodeCompose"),
386         NodeItem("TextureNodeDecompose"),
387         ]),
388     TextureNodeCategory("TEX_PATTERN", "Pattern", items=[
389         NodeItem("TextureNodeChecker"),
390         NodeItem("TextureNodeBricks"),
391         ]),
392     TextureNodeCategory("TEX_TEXTURE", "Textures", items=[
393         NodeItem("TextureNodeTexNoise"),
394         NodeItem("TextureNodeTexDistNoise"),
395         NodeItem("TextureNodeTexClouds"),
396         NodeItem("TextureNodeTexBlend"),
397         NodeItem("TextureNodeTexVoronoi"),
398         NodeItem("TextureNodeTexMagic"),
399         NodeItem("TextureNodeTexMarble"),
400         NodeItem("TextureNodeTexWood"),
401         NodeItem("TextureNodeTexMusgrave"),
402         NodeItem("TextureNodeTexStucci"),
403         ]),
404     TextureNodeCategory("TEX_CONVERTOR", "Converter", items=[
405         NodeItem("TextureNodeMath"),
406         NodeItem("TextureNodeValToRGB"),
407         NodeItem("TextureNodeRGBToBW"),
408         NodeItem("TextureNodeValToNor"),
409         NodeItem("TextureNodeDistance"),
410         ]),
411     TextureNodeCategory("TEX_DISTORT", "Distort", items=[
412         NodeItem("TextureNodeScale"),
413         NodeItem("TextureNodeTranslate"),
414         NodeItem("TextureNodeRotate"),
415         NodeItem("TextureNodeAt"),
416         ]),
417     TextureNodeCategory("TEX_GROUP", "Group", items=node_group_items),
418     TextureNodeCategory("TEX_LAYOUT", "Layout", items=[
419         NodeItem("NodeFrame"),
420         NodeItem("NodeReroute"),
421         ]),
422     ]
423
424
425 def register():
426     nodeitems_utils.register_node_categories('SHADER', shader_node_categories)
427     nodeitems_utils.register_node_categories('COMPOSITING', compositor_node_categories)
428     nodeitems_utils.register_node_categories('TEXTURE', texture_node_categories)
429
430
431 def unregister():
432     nodeitems_utils.unregister_node_categories('SHADER')
433     nodeitems_utils.unregister_node_categories('COMPOSITING')
434     nodeitems_utils.unregister_node_categories('TEXTURE')
435
436
437 if __name__ == "__main__":
438     register()