Merge branch 'master' into blender2.8
[blender.git] / release / scripts / modules / nodeitems_utils.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
22
23 class NodeCategory:
24     @classmethod
25     def poll(cls, context):
26         return True
27
28     def __init__(self, identifier, name, description="", items=None):
29         self.identifier = identifier
30         self.name = name
31         self.description = description
32
33         if items is None:
34             self.items = lambda context: []
35         elif callable(items):
36             self.items = items
37         else:
38             def items_gen(context):
39                 for item in items:
40                     if item.poll is None or context is None or item.poll(context):
41                         yield item
42             self.items = items_gen
43
44
45 class NodeItem:
46     def __init__(self, nodetype, label=None, settings=None, poll=None):
47
48         if settings is None:
49             settings = {}
50
51         self.nodetype = nodetype
52         self._label = label
53         self.settings = settings
54         self.poll = poll
55
56     @property
57     def label(self):
58         if self._label:
59             return self._label
60         else:
61             # if no custom label is defined, fall back to the node type UI name
62             cls = next(cls for cls in bpy.types.Node.__subclasses__() if cls.bl_rna.identifier == self.nodetype)
63             return cls.bl_rna.name
64
65     @property
66     def translation_context(self):
67         if self._label:
68             return bpy.app.translations.contexts.default
69         else:
70             # if no custom label is defined, fall back to the node type UI name
71             cls = next(cls for cls in bpy.types.Node.__subclasses__() if cls.bl_rna.identifier == self.nodetype)
72             return cls.bl_rna.translation_context
73
74     # NB: is a staticmethod because called with an explicit self argument
75     # NodeItemCustom sets this as a variable attribute in __init__
76     @staticmethod
77     def draw(self, layout, context):
78         props = layout.operator("node.add_node", text=self.label, text_ctxt=self.translation_context)
79         props.type = self.nodetype
80         props.use_transform = True
81
82         for setting in self.settings.items():
83             ops = props.settings.add()
84             ops.name = setting[0]
85             ops.value = setting[1]
86
87
88 class NodeItemCustom:
89     def __init__(self, poll=None, draw=None):
90         self.poll = poll
91         self.draw = draw
92
93
94 _node_categories = {}
95
96
97 def register_node_categories(identifier, cat_list):
98     if identifier in _node_categories:
99         raise KeyError("Node categories list '%s' already registered" % identifier)
100         return
101
102     # works as draw function for both menus and panels
103     def draw_node_item(self, context):
104         layout = self.layout
105         col = layout.column()
106         for item in self.category.items(context):
107             item.draw(item, col, context)
108
109     menu_types = []
110     panel_types = []
111     for cat in cat_list:
112         menu_type = type("NODE_MT_category_" + cat.identifier, (bpy.types.Menu,), {
113             "bl_space_type": 'NODE_EDITOR',
114             "bl_label": cat.name,
115             "category": cat,
116             "poll": cat.poll,
117             "draw": draw_node_item,
118             })
119         panel_type = type("NODE_PT_category_" + cat.identifier, (bpy.types.Panel,), {
120             "bl_space_type": 'NODE_EDITOR',
121             "bl_region_type": 'TOOLS',
122             "bl_label": cat.name,
123             "bl_category": cat.name,
124             "category": cat,
125             "poll": cat.poll,
126             "draw": draw_node_item,
127             })
128
129         menu_types.append(menu_type)
130         panel_types.append(panel_type)
131
132         bpy.utils.register_class(menu_type)
133         bpy.utils.register_class(panel_type)
134
135     def draw_add_menu(self, context):
136         layout = self.layout
137
138         for cat in cat_list:
139             if cat.poll(context):
140                 layout.menu("NODE_MT_category_%s" % cat.identifier)
141
142     # stores: (categories list, menu draw function, submenu types, panel types)
143     _node_categories[identifier] = (cat_list, draw_add_menu, menu_types, panel_types)
144
145
146 def node_categories_iter(context):
147     for cat_type in _node_categories.values():
148         for cat in cat_type[0]:
149             if cat.poll and ((context is None) or cat.poll(context)):
150                 yield cat
151
152
153 def node_items_iter(context):
154     for cat in node_categories_iter(context):
155         for item in cat.items(context):
156             yield item
157
158
159 def unregister_node_cat_types(cats):
160     for mt in cats[2]:
161         bpy.utils.unregister_class(mt)
162     for pt in cats[3]:
163         bpy.utils.unregister_class(pt)
164
165
166 def unregister_node_categories(identifier=None):
167     # unregister existing UI classes
168     if identifier:
169         cat_types = _node_categories.get(identifier, None)
170         if cat_types:
171             unregister_node_cat_types(cat_types)
172         del _node_categories[identifier]
173
174     else:
175         for cat_types in _node_categories.values():
176             unregister_node_cat_types(cat_types)
177         _node_categories.clear()
178
179
180 def draw_node_categories_menu(self, context):
181     for cats in _node_categories.values():
182         cats[1](self, context)