168 lines
4.8 KiB
Python
168 lines
4.8 KiB
Python
# SPDX-FileCopyrightText: 2013-2023 Blender Authors
|
|
#
|
|
# SPDX-License-Identifier: GPL-2.0-or-later
|
|
|
|
import bpy
|
|
|
|
|
|
class NodeCategory:
|
|
@classmethod
|
|
def poll(cls, _context):
|
|
return True
|
|
|
|
def __init__(self, identifier, name, *, description="", items=None):
|
|
self.identifier = identifier
|
|
self.name = name
|
|
self.description = description
|
|
|
|
if items is None:
|
|
self.items = lambda context: []
|
|
elif callable(items):
|
|
self.items = items
|
|
else:
|
|
def items_gen(context):
|
|
for item in items:
|
|
if item.poll is None or context is None or item.poll(context):
|
|
yield item
|
|
self.items = items_gen
|
|
|
|
|
|
class NodeItem:
|
|
def __init__(self, nodetype, *, label=None, settings=None, poll=None):
|
|
|
|
if settings is None:
|
|
settings = {}
|
|
|
|
self.nodetype = nodetype
|
|
self._label = label
|
|
self.settings = settings
|
|
self.poll = poll
|
|
|
|
@property
|
|
def label(self):
|
|
if self._label:
|
|
return self._label
|
|
else:
|
|
# if no custom label is defined, fall back to the node type UI name
|
|
bl_rna = bpy.types.Node.bl_rna_get_subclass(self.nodetype)
|
|
if bl_rna is not None:
|
|
return bl_rna.name
|
|
else:
|
|
return "Unknown"
|
|
|
|
@property
|
|
def translation_context(self):
|
|
if self._label:
|
|
return bpy.app.translations.contexts.default
|
|
else:
|
|
# if no custom label is defined, fall back to the node type UI name
|
|
bl_rna = bpy.types.Node.bl_rna_get_subclass(self.nodetype)
|
|
if bl_rna is not None:
|
|
return bl_rna.translation_context
|
|
else:
|
|
return bpy.app.translations.contexts.default
|
|
|
|
# NOTE: is a staticmethod because called with an explicit self argument
|
|
# NodeItemCustom sets this as a variable attribute in __init__
|
|
@staticmethod
|
|
def draw(self, layout, _context):
|
|
props = layout.operator("node.add_node", text=self.label, text_ctxt=self.translation_context)
|
|
props.type = self.nodetype
|
|
props.use_transform = True
|
|
|
|
for setting in self.settings.items():
|
|
ops = props.settings.add()
|
|
ops.name = setting[0]
|
|
ops.value = setting[1]
|
|
|
|
|
|
class NodeItemCustom:
|
|
def __init__(self, *, poll=None, draw=None):
|
|
self.poll = poll
|
|
self.draw = draw
|
|
|
|
|
|
_node_categories = {}
|
|
|
|
|
|
def register_node_categories(identifier, cat_list):
|
|
if identifier in _node_categories:
|
|
raise KeyError("Node categories list '%s' already registered" % identifier)
|
|
return
|
|
|
|
# works as draw function for menus
|
|
def draw_node_item(self, context):
|
|
layout = self.layout
|
|
col = layout.column(align=True)
|
|
for item in self.category.items(context):
|
|
item.draw(item, col, context)
|
|
|
|
menu_types = []
|
|
for cat in cat_list:
|
|
menu_type = type("NODE_MT_category_" + cat.identifier, (bpy.types.Menu,), {
|
|
"bl_space_type": 'NODE_EDITOR',
|
|
"bl_label": cat.name,
|
|
"category": cat,
|
|
"poll": cat.poll,
|
|
"draw": draw_node_item,
|
|
})
|
|
|
|
menu_types.append(menu_type)
|
|
|
|
bpy.utils.register_class(menu_type)
|
|
|
|
def draw_add_menu(self, context):
|
|
layout = self.layout
|
|
|
|
for cat in cat_list:
|
|
if cat.poll(context):
|
|
layout.menu("NODE_MT_category_%s" % cat.identifier)
|
|
|
|
# Stores: (categories list, menu draw function, sub-menu types).
|
|
_node_categories[identifier] = (cat_list, draw_add_menu, menu_types)
|
|
|
|
|
|
def node_categories_iter(context):
|
|
for cat_type in _node_categories.values():
|
|
for cat in cat_type[0]:
|
|
if cat.poll and ((context is None) or cat.poll(context)):
|
|
yield cat
|
|
|
|
|
|
def has_node_categories(context):
|
|
for cat_type in _node_categories.values():
|
|
for cat in cat_type[0]:
|
|
if cat.poll and ((context is None) or cat.poll(context)):
|
|
return True
|
|
return False
|
|
|
|
|
|
def node_items_iter(context):
|
|
for cat in node_categories_iter(context):
|
|
for item in cat.items(context):
|
|
yield item
|
|
|
|
|
|
def unregister_node_cat_types(cats):
|
|
for mt in cats[2]:
|
|
bpy.utils.unregister_class(mt)
|
|
|
|
|
|
def unregister_node_categories(identifier=None):
|
|
# unregister existing UI classes
|
|
if identifier:
|
|
cat_types = _node_categories.get(identifier, None)
|
|
if cat_types:
|
|
unregister_node_cat_types(cat_types)
|
|
del _node_categories[identifier]
|
|
|
|
else:
|
|
for cat_types in _node_categories.values():
|
|
unregister_node_cat_types(cat_types)
|
|
_node_categories.clear()
|
|
|
|
|
|
def draw_node_categories_menu(self, context):
|
|
for cats in _node_categories.values():
|
|
cats[1](self, context)
|