2022-11-12 23:59:00 +00:00
|
|
|
import bpy
|
2022-11-14 23:31:26 +00:00
|
|
|
import inspect
|
2022-11-12 23:59:00 +00:00
|
|
|
try:
|
|
|
|
import node_arrange as node_arrange
|
|
|
|
except:
|
|
|
|
pass
|
|
|
|
from .state import State
|
2022-11-13 19:38:31 +00:00
|
|
|
from .types import *
|
2022-11-12 23:59:00 +00:00
|
|
|
from .node_mapper import *
|
|
|
|
|
2022-11-13 19:38:31 +00:00
|
|
|
def _as_iterable(x):
|
2022-11-12 23:59:00 +00:00
|
|
|
try:
|
2022-11-13 19:38:31 +00:00
|
|
|
return iter(x)
|
2022-11-12 23:59:00 +00:00
|
|
|
except TypeError:
|
2022-11-13 19:38:31 +00:00
|
|
|
return [x,]
|
2022-11-12 23:59:00 +00:00
|
|
|
|
|
|
|
def tree(name):
|
|
|
|
tree_name = name
|
|
|
|
def build_tree(builder):
|
2022-11-14 23:31:26 +00:00
|
|
|
signature = inspect.signature(builder)
|
|
|
|
|
2022-11-12 23:59:00 +00:00
|
|
|
# Locate or create the node group
|
|
|
|
node_group = None
|
|
|
|
if tree_name in bpy.data.node_groups:
|
|
|
|
node_group = bpy.data.node_groups[tree_name]
|
|
|
|
else:
|
|
|
|
node_group = bpy.data.node_groups.new(tree_name, 'GeometryNodeTree')
|
|
|
|
# Clear the node group before building
|
|
|
|
for node in node_group.nodes:
|
|
|
|
node_group.nodes.remove(node)
|
2022-11-14 23:31:26 +00:00
|
|
|
while len(node_group.inputs) > len(signature.parameters):
|
|
|
|
node_group.inputs.remove(node_group.inputs[-1])
|
2022-11-12 23:59:00 +00:00
|
|
|
for group_output in node_group.outputs:
|
|
|
|
node_group.outputs.remove(group_output)
|
|
|
|
|
|
|
|
# Setup the group inputs
|
|
|
|
group_input_node = node_group.nodes.new('NodeGroupInput')
|
|
|
|
group_output_node = node_group.nodes.new('NodeGroupOutput')
|
|
|
|
|
|
|
|
# Collect the inputs
|
|
|
|
inputs = {}
|
2022-11-14 23:31:26 +00:00
|
|
|
for param in signature.parameters.values():
|
|
|
|
if param.annotation == inspect.Parameter.empty:
|
|
|
|
raise Exception(f"Tree input '{param.name}' has no type specified. Please annotate with a valid node input type.")
|
|
|
|
if not issubclass(param.annotation, Type):
|
|
|
|
raise Exception(f"Type of tree input '{param.name}' is not a valid 'Type' subclass.")
|
|
|
|
inputs[param.name] = (param.annotation, param.default)
|
2022-11-12 23:59:00 +00:00
|
|
|
|
|
|
|
# Create the input sockets and collect input values.
|
2022-11-14 23:31:26 +00:00
|
|
|
for i, node_input in enumerate(node_group.inputs):
|
|
|
|
if node_input.bl_socket_idname != list(inputs.values())[i][0].socket_type:
|
|
|
|
for ni in node_group.inputs:
|
|
|
|
node_group.inputs.remove(ni)
|
|
|
|
break
|
2022-11-12 23:59:00 +00:00
|
|
|
builder_inputs = []
|
|
|
|
for i, arg in enumerate(inputs.items()):
|
2022-11-14 23:31:26 +00:00
|
|
|
input_name = arg[0].replace('_', ' ').title()
|
|
|
|
if len(node_group.inputs) > i:
|
|
|
|
node_group.inputs[i].name = input_name
|
|
|
|
node_input = node_group.inputs[i]
|
|
|
|
else:
|
|
|
|
node_input = node_group.inputs.new(arg[1][0].socket_type, input_name)
|
|
|
|
if arg[1][1] != inspect.Parameter.empty:
|
|
|
|
node_input.default_value = arg[1][1]
|
|
|
|
builder_inputs.append(arg[1][0](group_input_node.outputs[i]))
|
2022-11-12 23:59:00 +00:00
|
|
|
|
|
|
|
# Run the builder function
|
|
|
|
State.current_node_tree = node_group
|
2022-11-15 03:32:32 +00:00
|
|
|
if inspect.isgeneratorfunction(builder):
|
|
|
|
generated_outputs = [*builder(*builder_inputs)]
|
|
|
|
if all(map(lambda x: issubclass(type(x), Type) and x._socket.type == 'GEOMETRY', generated_outputs)):
|
|
|
|
outputs = join_geometry(geometry=generated_outputs)
|
|
|
|
else:
|
|
|
|
outputs = generated_outputs
|
|
|
|
else:
|
|
|
|
outputs = builder(*builder_inputs)
|
2022-11-12 23:59:00 +00:00
|
|
|
|
|
|
|
# Create the output sockets
|
|
|
|
for i, result in enumerate(_as_iterable(outputs)):
|
|
|
|
if not issubclass(type(result), Type):
|
|
|
|
result = Type(value=result)
|
|
|
|
# raise Exception(f"Return value '{result}' is not a valid 'Type' subclass.")
|
|
|
|
node_group.outputs.new(result.socket_type, 'Result')
|
|
|
|
link = node_group.links.new(result._socket, group_output_node.inputs[i])
|
|
|
|
|
|
|
|
# Attempt to run the "Node Arrange" add-on on the tree.
|
|
|
|
try:
|
|
|
|
for area in bpy.context.screen.areas:
|
|
|
|
for space in area.spaces:
|
|
|
|
if space.type == 'NODE_EDITOR':
|
|
|
|
space.node_tree = node_group
|
|
|
|
with bpy.context.temp_override(area=area, space=space, space_data=space):
|
|
|
|
ntree = node_group
|
|
|
|
ntree.nodes[0].select = True
|
|
|
|
ntree.nodes.active = ntree.nodes[0]
|
|
|
|
n_groups = []
|
|
|
|
for i in ntree.nodes:
|
|
|
|
if i.type == 'GROUP':
|
|
|
|
n_groups.append(i)
|
|
|
|
|
|
|
|
while n_groups:
|
|
|
|
j = n_groups.pop(0)
|
|
|
|
node_arrange.nodes_iterate(j.node_tree)
|
|
|
|
for i in j.node_tree.nodes:
|
|
|
|
if i.type == 'GROUP':
|
|
|
|
n_groups.append(i)
|
|
|
|
|
|
|
|
node_arrange.nodes_iterate(ntree)
|
|
|
|
|
|
|
|
# arrange nodes + this center nodes together
|
|
|
|
if bpy.context.scene.node_center:
|
|
|
|
node_arrange.nodes_center(ntree)
|
|
|
|
except:
|
|
|
|
pass
|
|
|
|
|
|
|
|
# Return a function that creates a NodeGroup node in the tree.
|
|
|
|
# This lets @trees be used in other @trees via simple function calls.
|
|
|
|
def group_reference(*args, **kwargs):
|
|
|
|
return group(node_tree=node_group, *args, **kwargs)
|
|
|
|
return group_reference
|
|
|
|
if isinstance(name, str):
|
|
|
|
return build_tree
|
|
|
|
else:
|
|
|
|
tree_name = name.__name__
|
|
|
|
return build_tree(name)
|