250 lines
8.6 KiB
Python
250 lines
8.6 KiB
Python
# SPDX-License-Identifier: GPL-2.0-or-later
|
|
|
|
import bpy
|
|
from bpy.types import Operator
|
|
|
|
from bpy.app.translations import pgettext_data as data_
|
|
|
|
|
|
def build_default_empty_geometry_node_group(name):
|
|
group = bpy.data.node_groups.new(name, 'GeometryNodeTree')
|
|
group.inputs.new('NodeSocketGeometry', data_("Geometry"))
|
|
group.outputs.new('NodeSocketGeometry', data_("Geometry"))
|
|
input_node = group.nodes.new('NodeGroupInput')
|
|
output_node = group.nodes.new('NodeGroupOutput')
|
|
output_node.is_active_output = True
|
|
|
|
input_node.select = False
|
|
output_node.select = False
|
|
|
|
input_node.location.x = -200 - input_node.width
|
|
output_node.location.x = 200
|
|
|
|
return group
|
|
|
|
|
|
def geometry_node_group_empty_new():
|
|
group = build_default_empty_geometry_node_group(data_("Geometry Nodes"))
|
|
group.links.new(group.nodes[data_("Group Input")].outputs[0], group.nodes[data_("Group Output")].inputs[0])
|
|
return group
|
|
|
|
|
|
def geometry_modifier_poll(context):
|
|
ob = context.object
|
|
|
|
# Test object support for geometry node modifier
|
|
if not ob or ob.type not in {'MESH', 'POINTCLOUD', 'VOLUME', 'CURVE', 'FONT', 'CURVES'}:
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def get_context_modifier(context):
|
|
area = context.area
|
|
if (area is not None) and (area.type == 'PROPERTIES'):
|
|
modifier = context.modifier
|
|
else:
|
|
ob = context.object
|
|
if ob is None:
|
|
return False
|
|
modifier = ob.modifiers.active
|
|
if modifier is None or modifier.type != 'NODES':
|
|
return None
|
|
return modifier
|
|
|
|
|
|
def edit_geometry_nodes_modifier_poll(context):
|
|
return get_context_modifier(context) is not None
|
|
|
|
|
|
def socket_idname_to_attribute_type(idname):
|
|
if idname.startswith("NodeSocketInt"):
|
|
return "INT"
|
|
elif idname.startswith("NodeSocketColor"):
|
|
return "FLOAT_COLOR"
|
|
elif idname.startswith("NodeSocketVector"):
|
|
return "FLOAT_VECTOR"
|
|
elif idname.startswith("NodeSocketBool"):
|
|
return "BOOLEAN"
|
|
elif idname.startswith("NodeSocketFloat"):
|
|
return "FLOAT"
|
|
raise ValueError("Unsupported socket type")
|
|
return ""
|
|
|
|
|
|
def modifier_attribute_name_get(modifier, identifier):
|
|
try:
|
|
return modifier[identifier + "_attribute_name"]
|
|
except KeyError:
|
|
return None
|
|
|
|
|
|
def modifier_input_use_attribute(modifier, identifier):
|
|
try:
|
|
return modifier[identifier + "_use_attribute"] != 0
|
|
except KeyError:
|
|
return False
|
|
|
|
|
|
def get_socket_with_identifier(sockets, identifier):
|
|
for socket in sockets:
|
|
if socket.identifier == identifier:
|
|
return socket
|
|
return None
|
|
|
|
|
|
def get_enabled_socket_with_name(sockets, name):
|
|
for socket in sockets:
|
|
if socket.name == name and socket.enabled:
|
|
return socket
|
|
return None
|
|
|
|
|
|
class MoveModifierToNodes(Operator):
|
|
"""Move inputs and outputs from in the modifier to a new node group"""
|
|
|
|
bl_idname = "object.geometry_nodes_move_to_nodes"
|
|
bl_label = "Move to Nodes"
|
|
bl_options = {'REGISTER', 'UNDO'}
|
|
|
|
@classmethod
|
|
def poll(cls, context):
|
|
return edit_geometry_nodes_modifier_poll(context)
|
|
|
|
def execute(self, context):
|
|
modifier = get_context_modifier(context)
|
|
if not modifier:
|
|
return {'CANCELLED'}
|
|
old_group = modifier.node_group
|
|
if not old_group:
|
|
return {'CANCELLED'}
|
|
|
|
wrapper_name = old_group.name + ".wrapper"
|
|
group = build_default_empty_geometry_node_group(wrapper_name)
|
|
group_node = group.nodes.new("GeometryNodeGroup")
|
|
group_node.node_tree = old_group
|
|
group_node.update()
|
|
|
|
group_input_node = group.nodes[data_("Group Input")]
|
|
group_output_node = group.nodes[data_("Group Output")]
|
|
|
|
# Copy default values for inputs and create named attribute input nodes.
|
|
input_nodes = []
|
|
first_geometry_input = None
|
|
for input_socket in old_group.inputs:
|
|
identifier = input_socket.identifier
|
|
group_node_input = get_socket_with_identifier(group_node.inputs, identifier)
|
|
if modifier_input_use_attribute(modifier, identifier):
|
|
input_node = group.nodes.new("GeometryNodeInputNamedAttribute")
|
|
input_nodes.append(input_node)
|
|
input_node.data_type = socket_idname_to_attribute_type(input_socket.bl_socket_idname)
|
|
attribute_name = modifier_attribute_name_get(modifier, identifier)
|
|
input_node.inputs["Name"].default_value = attribute_name
|
|
output_socket = get_enabled_socket_with_name(input_node.outputs, "Attribute")
|
|
group.links.new(output_socket, group_node_input)
|
|
elif hasattr(input_socket, "default_value"):
|
|
group_node_input.default_value = modifier[identifier]
|
|
elif input_socket.bl_socket_idname == 'NodeSocketGeometry':
|
|
if not first_geometry_input:
|
|
first_geometry_input = group_node_input
|
|
|
|
group.links.new(group_input_node.outputs[0], first_geometry_input)
|
|
|
|
# Adjust locations of named attribute input nodes and group input node to make some space.
|
|
if input_nodes:
|
|
for i, node in enumerate(input_nodes):
|
|
node.location.x = -175
|
|
node.location.y = i * -50
|
|
group_input_node.location.x = -350
|
|
|
|
# Connect outputs to store named attribute nodes to replace modifier attribute outputs.
|
|
store_nodes = []
|
|
first_geometry_output = None
|
|
for output_socket in old_group.outputs:
|
|
identifier = output_socket.identifier
|
|
group_node_output = get_socket_with_identifier(group_node.outputs, identifier)
|
|
attribute_name = modifier_attribute_name_get(modifier, identifier)
|
|
if attribute_name:
|
|
store_node = group.nodes.new("GeometryNodeStoreNamedAttribute")
|
|
store_nodes.append(store_node)
|
|
store_node.data_type = socket_idname_to_attribute_type(output_socket.bl_socket_idname)
|
|
store_node.domain = output_socket.attribute_domain
|
|
store_node.inputs["Name"].default_value = attribute_name
|
|
input_socket = get_enabled_socket_with_name(store_node.inputs, "Value")
|
|
group.links.new(group_node_output, input_socket)
|
|
elif output_socket.bl_socket_idname == 'NodeSocketGeometry':
|
|
if not first_geometry_output:
|
|
first_geometry_output = group_node_output
|
|
|
|
# Adjust locations of store named attribute nodes and move group output.
|
|
# Note that the node group has its sockets names translated, while the built-in nodes don't.
|
|
if store_nodes:
|
|
for i, node in enumerate(store_nodes):
|
|
node.location.x = (i + 1) * 175
|
|
node.location.y = 0
|
|
group_output_node.location.x = (len(store_nodes) + 1) * 175
|
|
|
|
group.links.new(first_geometry_output, store_nodes[0].inputs["Geometry"])
|
|
for i in range(len(store_nodes) - 1):
|
|
group.links.new(store_nodes[i].outputs["Geometry"], store_nodes[i + 1].inputs["Geometry"])
|
|
|
|
group.links.new(store_nodes[-1].outputs["Geometry"], group_output_node.inputs[data_("Geometry")])
|
|
else:
|
|
group.links.new(first_geometry_output, group_output_node.inputs[data_("Geometry")])
|
|
|
|
modifier.node_group = group
|
|
|
|
return {'FINISHED'}
|
|
|
|
|
|
class NewGeometryNodesModifier(Operator):
|
|
"""Create a new modifier with a new geometry node group"""
|
|
|
|
bl_idname = "node.new_geometry_nodes_modifier"
|
|
bl_label = "New Geometry Node Modifier"
|
|
bl_options = {'REGISTER', 'UNDO'}
|
|
|
|
@classmethod
|
|
def poll(cls, context):
|
|
return geometry_modifier_poll(context)
|
|
|
|
def execute(self, context):
|
|
ob = context.object
|
|
modifier = ob.modifiers.new(data_("GeometryNodes"), 'NODES')
|
|
if not modifier:
|
|
return {'CANCELLED'}
|
|
|
|
group = geometry_node_group_empty_new()
|
|
modifier.node_group = group
|
|
|
|
return {'FINISHED'}
|
|
|
|
|
|
class NewGeometryNodeTreeAssign(Operator):
|
|
"""Create a new geometry node group and assign it to the active modifier"""
|
|
|
|
bl_idname = "node.new_geometry_node_group_assign"
|
|
bl_label = "Assign New Geometry Node Group"
|
|
bl_options = {'REGISTER', 'UNDO'}
|
|
|
|
@classmethod
|
|
def poll(cls, context):
|
|
return geometry_modifier_poll(context)
|
|
|
|
def execute(self, context):
|
|
modifier = get_context_modifier(context)
|
|
if not modifier:
|
|
return {'CANCELLED'}
|
|
|
|
group = geometry_node_group_empty_new()
|
|
modifier.node_group = group
|
|
|
|
return {'FINISHED'}
|
|
|
|
|
|
classes = (
|
|
NewGeometryNodesModifier,
|
|
NewGeometryNodeTreeAssign,
|
|
MoveModifierToNodes,
|
|
)
|