Workspace: Move engines to workspace and Properties Editor cleanup
[blender.git] / release / scripts / templates_py / custom_nodes.py
1 import bpy
2 from bpy.types import NodeTree, Node, NodeSocket
3
4 # Implementation of custom nodes from Python
5
6
7 # Derived from the NodeTree base type, similar to Menu, Operator, Panel, etc.
8 class MyCustomTree(NodeTree):
9     # Description string
10     '''A custom node tree type that will show up in the node editor header'''
11     # Optional identifier string. If not explicitly defined, the python class name is used.
12     bl_idname = 'CustomTreeType'
13     # Label for nice name display
14     bl_label = "Custom Node Tree"
15     # Icon identifier
16     bl_icon = 'NODETREE'
17
18
19 # Custom socket type
20 class MyCustomSocket(NodeSocket):
21     # Description string
22     '''Custom node socket type'''
23     # Optional identifier string. If not explicitly defined, the python class name is used.
24     bl_idname = 'CustomSocketType'
25     # Label for nice name display
26     bl_label = "Custom Node Socket"
27
28     # Enum items list
29     my_items = (
30         ('DOWN', "Down", "Where your feet are"),
31         ('UP', "Up", "Where your head should be"),
32         ('LEFT', "Left", "Not right"),
33         ('RIGHT', "Right", "Not left")
34     )
35
36     my_enum_prop = bpy.props.EnumProperty(name="Direction", description="Just an example", items=my_items, default='UP')
37
38     # Optional function for drawing the socket input value
39     def draw(self, context, layout, node, text):
40         if self.is_output or self.is_linked:
41             layout.label(text)
42         else:
43             layout.prop(self, "my_enum_prop", text=text)
44
45     # Socket color
46     def draw_color(self, context, node):
47         return (1.0, 0.4, 0.216, 0.5)
48
49
50 # Mix-in class for all custom nodes in this tree type.
51 # Defines a poll function to enable instantiation.
52 class MyCustomTreeNode:
53     @classmethod
54     def poll(cls, ntree):
55         return ntree.bl_idname == 'CustomTreeType'
56
57
58 # Derived from the Node base type.
59 class MyCustomNode(Node, MyCustomTreeNode):
60     # === Basics ===
61     # Description string
62     '''A custom node'''
63     # Optional identifier string. If not explicitly defined, the python class name is used.
64     bl_idname = 'CustomNodeType'
65     # Label for nice name display
66     bl_label = "Custom Node"
67     # Icon identifier
68     bl_icon = 'SOUND'
69
70     # === Custom Properties ===
71     # These work just like custom properties in ID data blocks
72     # Extensive information can be found under
73     # http://wiki.blender.org/index.php/Doc:2.6/Manual/Extensions/Python/Properties
74     my_string_prop = bpy.props.StringProperty()
75     my_float_prop = bpy.props.FloatProperty(default=3.1415926)
76
77     # === Optional Functions ===
78     # Initialization function, called when a new node is created.
79     # This is the most common place to create the sockets for a node, as shown below.
80     # NOTE: this is not the same as the standard __init__ function in Python, which is
81     #       a purely internal Python method and unknown to the node system!
82     def init(self, context):
83         self.inputs.new('CustomSocketType', "Hello")
84         self.inputs.new('NodeSocketFloat', "World")
85         self.inputs.new('NodeSocketVector', "!")
86
87         self.outputs.new('NodeSocketColor', "How")
88         self.outputs.new('NodeSocketColor', "are")
89         self.outputs.new('NodeSocketFloat', "you")
90
91     # Copy function to initialize a copied node from an existing one.
92     def copy(self, node):
93         print("Copying from node ", node)
94
95     # Free function to clean up on removal.
96     def free(self):
97         print("Removing node ", self, ", Goodbye!")
98
99     # Additional buttons displayed on the node.
100     def draw_buttons(self, context, layout):
101         layout.label("Node settings")
102         layout.prop(self, "my_float_prop")
103
104     # Detail buttons in the sidebar.
105     # If this function is not defined, the draw_buttons function is used instead
106     def draw_buttons_ext(self, context, layout):
107         layout.prop(self, "my_float_prop")
108         # my_string_prop button will only be visible in the sidebar
109         layout.prop(self, "my_string_prop")
110
111     # Optional: custom label
112     # Explicit user label overrides this, but here we can define a label dynamically
113     def draw_label(self):
114         return "I am a custom node"
115
116
117 ### Node Categories ###
118 # Node categories are a python system for automatically
119 # extending the Add menu, toolbar panels and search operator.
120 # For more examples see release/scripts/startup/nodeitems_builtins.py
121
122 import nodeitems_utils
123 from nodeitems_utils import NodeCategory, NodeItem
124
125 # our own base class with an appropriate poll function,
126 # so the categories only show up in our own tree type
127 class MyNodeCategory(NodeCategory):
128     @classmethod
129     def poll(cls, context):
130         return context.space_data.tree_type == 'CustomTreeType'
131
132
133 # all categories in a list
134 node_categories = [
135     # identifier, label, items list
136     MyNodeCategory('SOMENODES', "Some Nodes", items=[
137         # our basic node
138         NodeItem("CustomNodeType"),
139     ]),
140     MyNodeCategory('OTHERNODES', "Other Nodes", items=[
141         # the node item can have additional settings,
142         # which are applied to new nodes
143         # NB: settings values are stored as string expressions,
144         # for this reason they should be converted to strings using repr()
145         NodeItem("CustomNodeType", label="Node A", settings={
146             "my_string_prop": repr("Lorem ipsum dolor sit amet"),
147             "my_float_prop": repr(1.0),
148         }),
149         NodeItem("CustomNodeType", label="Node B", settings={
150             "my_string_prop": repr("consectetur adipisicing elit"),
151             "my_float_prop": repr(2.0),
152         }),
153     ]),
154 ]
155
156 classes = (
157     MyCustomTree,
158     MyCustomSocket,
159     MyCustomNode,
160 )
161
162 def register():
163     from bpy.utils import register_class
164     for cls in classes:
165         register_class(cls)
166
167     nodeitems_utils.register_node_categories('CUSTOM_NODES', node_categories)
168
169
170 def unregister():
171     nodeitems_utils.unregister_node_categories('CUSTOM_NODES')
172
173     from bpy.utils import unregister_class
174     for cls in reversed(classes):
175         unregister_class(cls)
176
177
178 if __name__ == "__main__":
179     register()