2020-09-07 00:06:41 +02:00
|
|
|
import itertools
|
2020-09-08 23:21:41 +02:00
|
|
|
from collections import OrderedDict
|
2020-09-06 19:06:28 +02:00
|
|
|
from typing import Any, Generator, List, Optional, Type, Dict
|
2020-09-07 00:06:41 +02:00
|
|
|
from typing import OrderedDict as ODict # Prevent naming conflicts
|
|
|
|
|
2017-04-08 00:34:45 +02:00
|
|
|
import bpy.types
|
|
|
|
from bpy.props import *
|
2017-03-06 02:29:03 +01:00
|
|
|
from nodeitems_utils import NodeItem
|
2020-09-07 00:06:41 +02:00
|
|
|
|
2020-09-07 19:15:39 +02:00
|
|
|
# When passed as a category to add_node(), this will use the capitalized
|
2020-09-09 20:48:35 +02:00
|
|
|
# name of the package of the node as the category to make renaming
|
|
|
|
# categories easier.
|
|
|
|
PKG_AS_CATEGORY = "__pkgcat__"
|
2020-09-07 19:15:39 +02:00
|
|
|
|
2017-03-06 02:29:03 +01:00
|
|
|
nodes = []
|
2020-09-07 00:06:41 +02:00
|
|
|
category_items: ODict[str, List['ArmNodeCategory']] = OrderedDict()
|
2017-04-08 00:34:45 +02:00
|
|
|
|
|
|
|
array_nodes = dict()
|
2017-03-06 02:29:03 +01:00
|
|
|
|
2020-09-07 01:13:00 +02:00
|
|
|
|
2020-09-08 21:16:20 +02:00
|
|
|
class ArmLogicTreeNode(bpy.types.Node):
|
2020-09-06 19:06:28 +02:00
|
|
|
def init(self, context):
|
|
|
|
# make sure a given node knows the version of the NodeClass from when it was created
|
|
|
|
if isinstance(type(self).arm_version, int):
|
|
|
|
self.arm_version = type(self).arm_version
|
|
|
|
else:
|
|
|
|
self.arm_version = 1
|
|
|
|
|
2017-03-06 02:29:03 +01:00
|
|
|
@classmethod
|
|
|
|
def poll(cls, ntree):
|
|
|
|
return ntree.bl_idname == 'ArmLogicTreeType'
|
|
|
|
|
2020-09-06 19:06:28 +02:00
|
|
|
def get_replacement_node(self, node_tree: bpy.types.NodeTree):
|
|
|
|
# needs to be overridden by individual node classes with arm_version>1
|
|
|
|
"""(only called if the node's version is inferior to the node class's version)
|
|
|
|
Help with the node replacement process, by explaining how a node (`self`) should be replaced.
|
|
|
|
This method can either return a NodeReplacement object (see `nodes_logic.py`), or a brand new node.
|
|
|
|
|
|
|
|
If a new node is returned, then the following needs to be already set:
|
|
|
|
- the node's links to the other nodes
|
|
|
|
- the node's properties
|
|
|
|
- the node inputs's default values
|
|
|
|
|
|
|
|
If more than one node need to be created (for example, if an input needs a type conversion after update),
|
|
|
|
please return all the nodes in a list.
|
|
|
|
|
|
|
|
please raise a LookupError specifically when the node's version isn't handled by the function.
|
|
|
|
|
|
|
|
note that the lowest 'defined' version should be 1. if the node's version is 0, it means that it has been saved before versioning was a thing.
|
|
|
|
NODES OF VERSION 1 AND VERSION 0 SHOULD HAVE THE SAME CONTENTS
|
|
|
|
"""
|
|
|
|
if self.arm_version==0 and type(self).arm_version == 1:
|
|
|
|
return NodeReplacement.Identity(self) # in case someone doesn't implement this function, but the node has version 0.
|
|
|
|
else:
|
|
|
|
raise LookupError(f"the current node class, {repr(type(self)):s}, does not implement the getReplacementNode method, even though it has updated")
|
|
|
|
|
2020-09-08 21:34:05 +02:00
|
|
|
def add_input(self, socket_type: str, socket_name: str, default_value: Any = None, is_var: bool = False) -> bpy.types.NodeSocket:
|
|
|
|
"""Adds a new input socket to the node.
|
|
|
|
|
|
|
|
If `is_var` is true, a dot is placed inside the socket to denote
|
|
|
|
that this socket can be used for variable access (see
|
|
|
|
SetVariable node).
|
|
|
|
"""
|
|
|
|
socket = self.inputs.new(socket_type, socket_name)
|
|
|
|
|
|
|
|
if default_value is not None:
|
|
|
|
socket.default_value = default_value
|
|
|
|
|
2020-09-08 22:08:29 +02:00
|
|
|
if is_var and not socket.display_shape.endswith('_DOT'):
|
|
|
|
socket.display_shape += '_DOT'
|
2020-09-08 21:34:05 +02:00
|
|
|
|
|
|
|
return socket
|
|
|
|
|
|
|
|
def add_output(self, socket_type: str, socket_name: str, default_value: Any = None, is_var: bool = False) -> bpy.types.NodeSocket:
|
|
|
|
"""Adds a new output socket to the node.
|
|
|
|
|
|
|
|
If `is_var` is true, a dot is placed inside the socket to denote
|
|
|
|
that this socket can be used for variable access (see
|
|
|
|
SetVariable node).
|
|
|
|
"""
|
|
|
|
socket = self.outputs.new(socket_type, socket_name)
|
|
|
|
|
|
|
|
if default_value is not None:
|
|
|
|
socket.default_value = default_value
|
|
|
|
|
|
|
|
if is_var and not socket.display_shape.endswith('_DOT'):
|
|
|
|
socket.display_shape += '_DOT'
|
|
|
|
|
|
|
|
return socket
|
|
|
|
|
2017-09-08 13:57:34 +02:00
|
|
|
|
2020-09-06 19:06:28 +02:00
|
|
|
class NodeReplacement:
|
|
|
|
"""
|
|
|
|
Represents a simple replacement rule, this can replace nodes of one type to nodes of a second type.
|
|
|
|
However, it is fairly limited. For instance, it assumes there are no changes in the type of the inputs or outputs
|
|
|
|
Second, it also assumes that node properties (especially EnumProperties) keep the same possible values.
|
|
|
|
|
|
|
|
- from_node, from_node_version: the type of node to be removed, and its version number
|
|
|
|
- to_node, to_node_version: the type of node which takes from_node's place, and its version number
|
|
|
|
- *_socket_mapping: a map which defines how the sockets of the old node shall be connected to the new node
|
|
|
|
{1: 2} means that anything connected to the socket with index 1 on the original node will be connected to the socket with index 2 on the new node
|
|
|
|
- property_mapping: the mapping used to transfer the values of the old node's properties to the new node's properties.
|
|
|
|
{"property0": "property1"} mean that the value of the new node's property1 should be the old node's property0's value.
|
|
|
|
- input_defaults: a mapping used to give default values to the inputs which aren't overridden otherwise.
|
|
|
|
- property_defaults: a mapping used to define the value of the new node's properties, when they aren't overridden otherwise.
|
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(self, from_node: str, from_node_version: int, to_node: str, to_node_version: int,
|
2020-09-13 15:32:39 +02:00
|
|
|
in_socket_mapping: Dict[int, int], out_socket_mapping: Dict[int, int], property_mapping: Optional[Dict[str, str]] = None,
|
|
|
|
input_defaults: Optional[Dict[int, any]] = None, property_defaults: Optional[Dict[str, any]] = None):
|
2020-09-06 19:06:28 +02:00
|
|
|
self.from_node = from_node
|
|
|
|
self.to_node = to_node
|
|
|
|
self.from_node_version = from_node_version
|
|
|
|
self.to_node_version = to_node_version
|
|
|
|
|
|
|
|
self.in_socket_mapping = in_socket_mapping
|
|
|
|
self.out_socket_mapping = out_socket_mapping
|
2020-09-13 15:32:39 +02:00
|
|
|
self.property_mapping = {} if property_mapping is None else property_mapping
|
2020-09-06 19:06:28 +02:00
|
|
|
|
2020-09-13 15:32:39 +02:00
|
|
|
self.input_defaults = {} if input_defaults is None else input_defaults
|
|
|
|
self.property_defaults = {} if property_defaults is None else property_defaults
|
2020-09-06 19:06:28 +02:00
|
|
|
|
|
|
|
@classmethod
|
|
|
|
def Identity(cls, node: ArmLogicTreeNode):
|
|
|
|
"""returns a NodeReplacement that does nothing, while operating on a given node.
|
|
|
|
WARNING: it assumes that all node properties are called "property0", "property1", etc...
|
|
|
|
"""
|
|
|
|
in_socks = {i:i for i in range(len(node.inputs))}
|
|
|
|
out_socks = {i:i for i in range(len(node.outputs))}
|
|
|
|
props = {}
|
|
|
|
i=0
|
|
|
|
while hasattr(node, f'property{i:d}'):
|
|
|
|
props[f'property{i:d}'] = f'property{i:d}'
|
|
|
|
i +=1
|
|
|
|
return NodeReplacement(
|
|
|
|
node.bl_idname, node.arm_version, node.bl_idname, type(node).arm_version,
|
|
|
|
in_socket_mapping=in_socks, out_socket_mapping=out_socks,
|
|
|
|
property_mapping=props
|
|
|
|
)
|
|
|
|
|
|
|
|
def chain_with(self, other):
|
|
|
|
"""modify the current NodeReplacement by "adding" a second replacement after it"""
|
|
|
|
if self.to_node != other.from_node or self.to_node_version != other.from_node_version:
|
|
|
|
raise TypeError('the given NodeReplacement-s could not be chained')
|
|
|
|
self.to_node = other.to_node
|
|
|
|
self.to_node_version = other.to_node_version
|
|
|
|
|
|
|
|
for i1, i2 in self.in_socket_mapping.items():
|
|
|
|
i3 = other.in_socket_mapping[i2]
|
|
|
|
self.in_socket_mapping[i1] = i3
|
|
|
|
for i1, i2 in self.out_socket_mapping.items():
|
|
|
|
i3 = other.out_socket_mapping[i2]
|
|
|
|
self.out_socket_mapping[i1] = i3
|
|
|
|
for p1, p2 in self.property_mapping.items():
|
|
|
|
p3 = other.property_mapping[p2]
|
|
|
|
self.property_mapping[p1] = p3
|
|
|
|
|
|
|
|
old_input_defaults = self.input_defaults
|
|
|
|
self.input_defaults = other.input_defaults.copy()
|
|
|
|
for i, x in old_input_defaults.items():
|
|
|
|
self.input_defaults[ other.in_socket_mapping[i] ] = x
|
|
|
|
|
|
|
|
old_property_defaults = self.property_defaults
|
|
|
|
self.property_defaults = other.property_defaults.copy()
|
|
|
|
for p, x in old_property_defaults.items():
|
|
|
|
self.property_defaults[ other.property_mapping[p] ] = x
|
|
|
|
|
2017-04-08 00:34:45 +02:00
|
|
|
class ArmNodeAddInputButton(bpy.types.Operator):
|
2020-09-14 23:20:19 +02:00
|
|
|
"""Add a new input socket to the node set by node_index."""
|
2017-04-08 00:34:45 +02:00
|
|
|
bl_idname = 'arm.node_add_input'
|
|
|
|
bl_label = 'Add Input'
|
2020-09-14 23:20:19 +02:00
|
|
|
|
2018-12-18 23:48:38 +01:00
|
|
|
node_index: StringProperty(name='Node Index', default='')
|
|
|
|
socket_type: StringProperty(name='Socket Type', default='NodeSocketShader')
|
|
|
|
name_format: StringProperty(name='Name Format', default='Input {0}')
|
|
|
|
index_name_offset: IntProperty(name='Index Name Offset', default=0)
|
2017-04-08 00:34:45 +02:00
|
|
|
|
|
|
|
def execute(self, context):
|
|
|
|
global array_nodes
|
2017-04-11 11:28:22 +02:00
|
|
|
inps = array_nodes[self.node_index].inputs
|
2018-11-04 05:05:52 +01:00
|
|
|
inps.new(self.socket_type, self.name_format.format(str(len(inps) + self.index_name_offset)))
|
2020-09-14 23:20:19 +02:00
|
|
|
|
|
|
|
# Reset to default again for subsequent calls of this operator
|
|
|
|
self.node_index = ''
|
|
|
|
self.socket_type = 'NodeSocketShader'
|
|
|
|
self.name_format = 'Input {0}'
|
|
|
|
self.index_name_offset = 0
|
|
|
|
|
2017-04-08 00:34:45 +02:00
|
|
|
return{'FINISHED'}
|
|
|
|
|
2018-06-12 23:48:01 +02:00
|
|
|
class ArmNodeAddInputValueButton(bpy.types.Operator):
|
2020-09-07 01:15:12 +02:00
|
|
|
"""Add new input"""
|
2018-06-12 23:48:01 +02:00
|
|
|
bl_idname = 'arm.node_add_input_value'
|
|
|
|
bl_label = 'Add Input'
|
2018-12-18 23:48:38 +01:00
|
|
|
node_index: StringProperty(name='Node Index', default='')
|
|
|
|
socket_type: StringProperty(name='Socket Type', default='NodeSocketShader')
|
2018-06-12 23:48:01 +02:00
|
|
|
|
|
|
|
def execute(self, context):
|
|
|
|
global array_nodes
|
|
|
|
inps = array_nodes[self.node_index].inputs
|
|
|
|
inps.new(self.socket_type, 'Value')
|
|
|
|
return{'FINISHED'}
|
|
|
|
|
2017-04-08 00:34:45 +02:00
|
|
|
class ArmNodeRemoveInputButton(bpy.types.Operator):
|
2020-09-07 01:15:12 +02:00
|
|
|
"""Remove last input"""
|
2017-04-08 00:34:45 +02:00
|
|
|
bl_idname = 'arm.node_remove_input'
|
|
|
|
bl_label = 'Remove Input'
|
2018-12-18 23:48:38 +01:00
|
|
|
node_index: StringProperty(name='Node Index', default='')
|
2017-04-08 00:34:45 +02:00
|
|
|
|
|
|
|
def execute(self, context):
|
|
|
|
global array_nodes
|
2018-02-26 03:43:06 +01:00
|
|
|
node = array_nodes[self.node_index]
|
|
|
|
inps = node.inputs
|
|
|
|
min_inps = 0 if not hasattr(node, 'min_inputs') else node.min_inputs
|
|
|
|
if len(inps) > min_inps:
|
2017-04-08 00:34:45 +02:00
|
|
|
inps.remove(inps.values()[-1])
|
|
|
|
return{'FINISHED'}
|
|
|
|
|
2018-06-12 23:48:01 +02:00
|
|
|
class ArmNodeRemoveInputValueButton(bpy.types.Operator):
|
2020-09-07 01:15:12 +02:00
|
|
|
"""Remove last input"""
|
2018-06-12 23:48:01 +02:00
|
|
|
bl_idname = 'arm.node_remove_input_value'
|
|
|
|
bl_label = 'Remove Input'
|
2018-12-18 23:48:38 +01:00
|
|
|
node_index: StringProperty(name='Node Index', default='')
|
2018-06-12 23:48:01 +02:00
|
|
|
|
|
|
|
def execute(self, context):
|
|
|
|
global array_nodes
|
|
|
|
node = array_nodes[self.node_index]
|
|
|
|
inps = node.inputs
|
|
|
|
min_inps = 0 if not hasattr(node, 'min_inputs') else node.min_inputs
|
|
|
|
if len(inps) > min_inps and inps[-1].name == 'Value':
|
|
|
|
inps.remove(inps.values()[-1])
|
|
|
|
return{'FINISHED'}
|
|
|
|
|
2017-04-11 11:28:22 +02:00
|
|
|
class ArmNodeAddOutputButton(bpy.types.Operator):
|
2020-09-14 23:20:19 +02:00
|
|
|
"""Add a new output socket to the node set by node_index"""
|
2017-04-11 11:28:22 +02:00
|
|
|
bl_idname = 'arm.node_add_output'
|
|
|
|
bl_label = 'Add Output'
|
2020-09-14 23:20:19 +02:00
|
|
|
|
2018-12-18 23:48:38 +01:00
|
|
|
node_index: StringProperty(name='Node Index', default='')
|
|
|
|
socket_type: StringProperty(name='Socket Type', default='NodeSocketShader')
|
|
|
|
name_format: StringProperty(name='Name Format', default='Output {0}')
|
|
|
|
index_name_offset: IntProperty(name='Index Name Offset', default=0)
|
2017-04-11 11:28:22 +02:00
|
|
|
|
|
|
|
def execute(self, context):
|
|
|
|
global array_nodes
|
|
|
|
outs = array_nodes[self.node_index].outputs
|
2018-11-04 05:05:52 +01:00
|
|
|
outs.new(self.socket_type, self.name_format.format(str(len(outs) + self.index_name_offset)))
|
2020-09-14 23:20:19 +02:00
|
|
|
|
|
|
|
# Reset to default again for subsequent calls of this operator
|
|
|
|
self.node_index = ''
|
|
|
|
self.socket_type = 'NodeSocketShader'
|
|
|
|
self.name_format = 'Output {0}'
|
|
|
|
self.index_name_offset = 0
|
|
|
|
|
2017-04-11 11:28:22 +02:00
|
|
|
return{'FINISHED'}
|
|
|
|
|
|
|
|
class ArmNodeRemoveOutputButton(bpy.types.Operator):
|
2020-09-07 01:15:12 +02:00
|
|
|
"""Remove last output"""
|
2017-04-11 11:28:22 +02:00
|
|
|
bl_idname = 'arm.node_remove_output'
|
|
|
|
bl_label = 'Remove Output'
|
2018-12-18 23:48:38 +01:00
|
|
|
node_index: StringProperty(name='Node Index', default='')
|
2017-04-11 11:28:22 +02:00
|
|
|
|
|
|
|
def execute(self, context):
|
|
|
|
global array_nodes
|
2018-08-12 08:42:39 +02:00
|
|
|
node = array_nodes[self.node_index]
|
|
|
|
outs = node.outputs
|
|
|
|
min_outs = 0 if not hasattr(node, 'min_outputs') else node.min_outputs
|
|
|
|
if len(outs) > min_outs:
|
|
|
|
outs.remove(outs.values()[-1])
|
|
|
|
return{'FINISHED'}
|
|
|
|
|
|
|
|
class ArmNodeAddInputOutputButton(bpy.types.Operator):
|
2020-09-07 01:15:12 +02:00
|
|
|
"""Add new input and output"""
|
2018-08-12 08:42:39 +02:00
|
|
|
bl_idname = 'arm.node_add_input_output'
|
|
|
|
bl_label = 'Add Input Output'
|
2020-09-14 23:20:19 +02:00
|
|
|
|
2018-12-18 23:48:38 +01:00
|
|
|
node_index: StringProperty(name='Node Index', default='')
|
|
|
|
in_socket_type: StringProperty(name='In Socket Type', default='NodeSocketShader')
|
|
|
|
out_socket_type: StringProperty(name='Out Socket Type', default='NodeSocketShader')
|
|
|
|
in_name_format: StringProperty(name='In Name Format', default='Input {0}')
|
|
|
|
out_name_format: StringProperty(name='Out Name Format', default='Output {0}')
|
|
|
|
in_index_name_offset: IntProperty(name='Index Name Offset', default=0)
|
2018-08-12 08:42:39 +02:00
|
|
|
|
|
|
|
def execute(self, context):
|
|
|
|
global array_nodes
|
|
|
|
node = array_nodes[self.node_index]
|
|
|
|
inps = node.inputs
|
|
|
|
outs = node.outputs
|
|
|
|
inps.new(self.in_socket_type, self.in_name_format.format(str(len(inps) + self.in_index_name_offset)))
|
|
|
|
outs.new(self.out_socket_type, self.out_name_format.format(str(len(outs))))
|
2020-09-14 23:20:19 +02:00
|
|
|
|
|
|
|
# Reset to default again for subsequent calls of this operator
|
|
|
|
self.node_index = ''
|
|
|
|
self.in_socket_type = 'NodeSocketShader'
|
|
|
|
self.out_socket_type = 'NodeSocketShader'
|
|
|
|
self.in_name_format = 'Input {0}'
|
|
|
|
self.out_name_format = 'Output {0}'
|
|
|
|
self.in_index_name_offset = 0
|
|
|
|
|
2018-08-12 08:42:39 +02:00
|
|
|
return{'FINISHED'}
|
|
|
|
|
|
|
|
class ArmNodeRemoveInputOutputButton(bpy.types.Operator):
|
2020-09-07 01:15:12 +02:00
|
|
|
"""Remove last input and output"""
|
2018-08-12 08:42:39 +02:00
|
|
|
bl_idname = 'arm.node_remove_input_output'
|
|
|
|
bl_label = 'Remove Input Output'
|
2018-12-18 23:48:38 +01:00
|
|
|
node_index: StringProperty(name='Node Index', default='')
|
2018-08-12 08:42:39 +02:00
|
|
|
|
|
|
|
def execute(self, context):
|
|
|
|
global array_nodes
|
|
|
|
node = array_nodes[self.node_index]
|
|
|
|
inps = node.inputs
|
|
|
|
outs = node.outputs
|
|
|
|
min_inps = 0 if not hasattr(node, 'min_inputs') else node.min_inputs
|
|
|
|
min_outs = 0 if not hasattr(node, 'min_outputs') else node.min_outputs
|
|
|
|
if len(inps) > min_inps:
|
|
|
|
inps.remove(inps.values()[-1])
|
|
|
|
if len(outs) > min_outs:
|
2017-04-11 11:28:22 +02:00
|
|
|
outs.remove(outs.values()[-1])
|
|
|
|
return{'FINISHED'}
|
|
|
|
|
2020-09-07 00:06:41 +02:00
|
|
|
|
2020-09-07 00:54:52 +02:00
|
|
|
class ArmNodeSearch(bpy.types.Operator):
|
|
|
|
bl_idname = "arm.node_search"
|
|
|
|
bl_label = "Search..."
|
|
|
|
bl_options = {"REGISTER"}
|
|
|
|
bl_property = "item"
|
|
|
|
|
|
|
|
def get_search_items(self, context):
|
|
|
|
items = []
|
|
|
|
for node in get_all_nodes():
|
|
|
|
items.append((node.nodetype, node.label, ""))
|
|
|
|
return items
|
|
|
|
|
|
|
|
item: EnumProperty(items=get_search_items)
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
def poll(cls, context):
|
|
|
|
return context.space_data.tree_type == 'ArmLogicTreeType' and context.space_data.edit_tree
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
def description(cls, context, properties):
|
|
|
|
if cls.poll(context):
|
|
|
|
return "Search for a logic node"
|
|
|
|
else:
|
|
|
|
return "Search for a logic node. This operator is not available" \
|
|
|
|
" without an active node tree"
|
|
|
|
|
|
|
|
def invoke(self, context, event):
|
|
|
|
context.window_manager.invoke_search_popup(self)
|
|
|
|
return {"CANCELLED"}
|
|
|
|
|
|
|
|
def execute(self, context):
|
|
|
|
"""Called when a node is added."""
|
|
|
|
bpy.ops.node.add_node('INVOKE_DEFAULT', type=self.item, use_transform=True)
|
|
|
|
return {"FINISHED"}
|
|
|
|
|
2020-09-07 00:06:41 +02:00
|
|
|
|
|
|
|
class ArmNodeCategory:
|
|
|
|
"""Represents a category (=directory) of logic nodes."""
|
|
|
|
def __init__(self, name: str, icon: str, description: str):
|
|
|
|
self.name = name
|
|
|
|
self.icon = icon
|
|
|
|
self.description = description
|
2020-09-08 01:26:41 +02:00
|
|
|
self.node_sections: ODict[str, List[NodeItem]] = OrderedDict()
|
2020-09-07 00:06:41 +02:00
|
|
|
|
2020-09-07 00:33:31 +02:00
|
|
|
def register_node(self, node_type: Type[bpy.types.Node], node_section: str) -> None:
|
2020-09-07 00:06:41 +02:00
|
|
|
"""Registers a node to this category so that it will be
|
|
|
|
displayed int the `Add node` menu."""
|
|
|
|
self.add_node_section(node_section)
|
|
|
|
|
2020-09-07 00:33:31 +02:00
|
|
|
# Internal node types seem to have no bl_idname attribute
|
|
|
|
if issubclass(node_type, bpy.types.NodeInternal):
|
|
|
|
item = NodeItem(node_type.__name__)
|
|
|
|
else:
|
|
|
|
item = NodeItem(node_type.bl_idname)
|
|
|
|
|
|
|
|
self.node_sections[node_section].append(item)
|
2020-09-07 00:06:41 +02:00
|
|
|
|
|
|
|
def get_all_nodes(self) -> Generator[NodeItem, None, None]:
|
|
|
|
"""Returns all nodes that are registered into this category."""
|
|
|
|
yield from itertools.chain(*self.node_sections.values())
|
|
|
|
|
|
|
|
def add_node_section(self, name: str):
|
|
|
|
"""Adds a node section to this category."""
|
|
|
|
if name not in self.node_sections:
|
|
|
|
self.node_sections[name] = []
|
|
|
|
|
2020-09-07 20:36:30 +02:00
|
|
|
def sort_nodes(self):
|
|
|
|
for node_section in self.node_sections:
|
|
|
|
self.node_sections[node_section] = sorted(self.node_sections[node_section], key=lambda item: item.label)
|
|
|
|
|
2020-09-07 00:06:41 +02:00
|
|
|
|
|
|
|
def category_exists(name: str) -> bool:
|
|
|
|
for category_section in category_items:
|
|
|
|
for c in category_items[category_section]:
|
|
|
|
if c.name == name:
|
|
|
|
return True
|
|
|
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
|
|
def get_category(name: str) -> Optional[ArmNodeCategory]:
|
|
|
|
for category_section in category_items:
|
|
|
|
for c in category_items[category_section]:
|
|
|
|
if c.name == name:
|
|
|
|
return c
|
|
|
|
|
|
|
|
return None
|
|
|
|
|
|
|
|
|
|
|
|
def get_all_categories() -> Generator[ArmNodeCategory, None, None]:
|
|
|
|
for section_categories in category_items.values():
|
|
|
|
yield from itertools.chain(section_categories)
|
|
|
|
|
|
|
|
|
|
|
|
def get_all_nodes() -> Generator[NodeItem, None, None]:
|
|
|
|
for category in get_all_categories():
|
|
|
|
yield from itertools.chain(category.get_all_nodes())
|
|
|
|
|
|
|
|
|
|
|
|
def add_category_section(name: str) -> None:
|
|
|
|
"""Adds a section of categories to the node menu to group multiple
|
|
|
|
categories visually together. The given name only acts as an ID and
|
|
|
|
is not displayed in the user inferface."""
|
|
|
|
global category_items
|
|
|
|
if name not in category_items:
|
|
|
|
category_items[name] = []
|
|
|
|
|
|
|
|
|
|
|
|
def add_node_section(name: str, category: str) -> None:
|
|
|
|
"""Adds a section of nodes to the sub menu of the given category to
|
|
|
|
group multiple nodes visually together. The given name only acts as
|
|
|
|
an ID and is not displayed in the user inferface."""
|
|
|
|
node_category = get_category(category)
|
|
|
|
|
|
|
|
if node_category is not None:
|
|
|
|
node_category.add_node_section(name)
|
|
|
|
|
|
|
|
|
|
|
|
def add_category(category: str, section: str = 'default', icon: str = 'BLANK1', description: str = '') -> Optional[ArmNodeCategory]:
|
|
|
|
"""Adds a category of nodes to the node menu."""
|
|
|
|
global category_items
|
|
|
|
|
|
|
|
add_category_section(section)
|
|
|
|
if not category_exists(category):
|
|
|
|
node_category = ArmNodeCategory(category, icon, description)
|
|
|
|
category_items[section].append(node_category)
|
|
|
|
return node_category
|
|
|
|
|
|
|
|
return None
|
|
|
|
|
|
|
|
|
2020-09-06 19:06:28 +02:00
|
|
|
def add_node(node_type: Type[bpy.types.Node], category: str, section: str = 'default', is_obselete: bool = False) -> None:
|
2020-09-07 00:06:41 +02:00
|
|
|
"""
|
|
|
|
Registers a node to the given category. If no section is given, the
|
|
|
|
node is put into the default section that does always exist.
|
|
|
|
"""
|
2017-04-08 00:34:45 +02:00
|
|
|
global nodes
|
2020-09-07 00:06:41 +02:00
|
|
|
|
2020-09-09 20:48:35 +02:00
|
|
|
if category == PKG_AS_CATEGORY:
|
2020-09-07 19:15:39 +02:00
|
|
|
category = node_type.__module__.rsplit('.', 2)[1].capitalize()
|
|
|
|
|
2020-09-07 00:33:16 +02:00
|
|
|
nodes.append(node_type)
|
2020-09-07 00:06:41 +02:00
|
|
|
node_category = get_category(category)
|
|
|
|
|
2020-09-06 19:06:28 +02:00
|
|
|
if is_obselete:
|
|
|
|
# We need the obselete nodes to be registered in order to have them replaced,
|
|
|
|
# but do not add them to the menu.
|
|
|
|
return
|
|
|
|
|
2020-09-07 00:06:41 +02:00
|
|
|
if node_category is None:
|
|
|
|
node_category = add_category(category)
|
|
|
|
|
2020-09-07 00:33:16 +02:00
|
|
|
node_category.register_node(node_type, section)
|
|
|
|
node_type.bl_icon = node_category.icon
|
2017-04-08 00:34:45 +02:00
|
|
|
|
2020-09-07 00:54:52 +02:00
|
|
|
|
|
|
|
bpy.utils.register_class(ArmNodeSearch)
|
2017-04-08 00:34:45 +02:00
|
|
|
bpy.utils.register_class(ArmNodeAddInputButton)
|
2018-06-12 23:48:01 +02:00
|
|
|
bpy.utils.register_class(ArmNodeAddInputValueButton)
|
2017-04-08 00:34:45 +02:00
|
|
|
bpy.utils.register_class(ArmNodeRemoveInputButton)
|
2018-06-12 23:48:01 +02:00
|
|
|
bpy.utils.register_class(ArmNodeRemoveInputValueButton)
|
2017-04-11 11:28:22 +02:00
|
|
|
bpy.utils.register_class(ArmNodeAddOutputButton)
|
|
|
|
bpy.utils.register_class(ArmNodeRemoveOutputButton)
|
2018-08-12 08:42:39 +02:00
|
|
|
bpy.utils.register_class(ArmNodeAddInputOutputButton)
|
|
|
|
bpy.utils.register_class(ArmNodeRemoveInputOutputButton)
|