Files
test2/scripts/startup/bl_operators/geometry_nodes.py
Damien Picard ae7c71ef09 Fix Python error when trying to call the Search Menu with no object
When pressing F3 to bring up the Search Menu, an error appeared
because the MoveModifierToNodes operator's poll method did not
consider the context where no object was selected.

Pull Request: https://projects.blender.org/blender/blender/pulls/107310
2023-04-24 20:47:20 +02:00

247 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):
if context.area.type == 'PROPERTIES':
modifier = context.modifier
else:
if context.object is None:
return False
modifier = context.object.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):
modifier = context.object.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,
)