test2/scripts/startup/bl_operators/geometry_nodes.py

# 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:
        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,
)