geometry-script/operators/convert_tree.py

import collections.abc
import keyword
import re

import bpy
import mathutils


_SOCKET_TYPE_NAMES = {
    "BOOLEAN": "Bool",
    "COLLECTION": "Collection",
    "GEOMETRY": "Geometry",
    "IMAGE": "Image",
    "INT": "Int",
    "MATERIAL": "Material",
    "MENU": "Menu",
    "OBJECT": "Object",
    "RGBA": "Color",
    "ROTATION": "Rotation",
    "STRING": "String",
    "TEXTURE": "Texture",
    "VALUE": "Float",
    "VECTOR": "Vector",
}

_PROP_DENYLIST = {
    "bl_idname",
    "bl_label",
    "bl_static_type",
    "color",
    "dimensions",
    "height",
    "hide",
    "inputs",
    "internal_links",
    "label",
    "location",
    "mute",
    "name",
    "outputs",
    "parent",
    "rna_type",
    "select",
    "show_options",
    "show_preview",
    "show_texture",
    "type",
    "use_custom_color",
    "warning_propagation",
    "width",
}


def _snake_case(value):
    value = re.sub(r"[^0-9A-Za-z_]+", "_", value.strip())
    value = re.sub(r"_+", "_", value).strip("_").lower()
    if not value:
        value = "value"
    if value[0].isdigit():
        value = "_" + value
    if keyword.iskeyword(value):
        value += "_"
    return value


def _node_function_name(node):
    if node.bl_idname == "GeometryNodeGroup" and getattr(node, "node_tree", None):
        return _snake_case(node.node_tree.name)
    return _snake_case(node.bl_rna.name)


def _socket_name(socket):
    return _snake_case(socket.name or socket.identifier)


def _enabled_sockets(sockets):
    return [socket for socket in sockets if socket.enabled and socket.type != "CUSTOM"]


def _convert_value(value):
    if isinstance(value, mathutils.Vector):
        return tuple(value)
    if isinstance(value, mathutils.Euler):
        return tuple(value)
    if isinstance(value, mathutils.Color):
        return tuple(value)
    if isinstance(value, bpy.types.bpy_prop_array):
        return tuple(value)
    if isinstance(value, collections.abc.Iterable) and not isinstance(value, (str, bytes)):
        try:
            return tuple(value)
        except TypeError:
            return value
    return value


def _literal(value):
    value = _convert_value(value)
    if isinstance(value, bpy.types.ID):
        return "None"
    return repr(value)


def _socket_type_annotation(socket):
    return _SOCKET_TYPE_NAMES.get(socket.type, "Type")


class _Assignment:
    def __init__(self, name, node):
        self.name = name
        self.node = node
        self.props = {}
        self.arguments = {}
        self.argument_outputs = {}
        self.argument_is_multi_input = {}

    def flattened_arguments(self):
        result = []
        for value in self.arguments.values():
            if isinstance(value, _Assignment):
                result.append(value)
            elif isinstance(value, list):
                result.extend(item for item in value if isinstance(item, _Assignment))
        return result

    def _value_expr(self, key, value, index=None):
        if isinstance(value, list):
            values = [self._value_expr(key, item, index=i) for i, item in enumerate(value)]
            if self.argument_is_multi_input.get(key, False):
                return "[" + ", ".join(values) + "]"
            if len(values) == 1:
                return values[0]
            return "(" + ", ".join(values) + ")"
        if not isinstance(value, _Assignment):
            return _literal(value)
        if value.node.type == "GROUP_INPUT":
            output_name = self.argument_outputs[key] if index is None else self.argument_outputs[key][index]
            return output_name
        output_name = self.argument_outputs[key] if index is None else self.argument_outputs[key][index]
        enabled_outputs = _enabled_sockets(value.node.outputs)
        if len(enabled_outputs) > 1:
            return f"{value.name}.{output_name}"
        return value.name

    def _argument_expr(self, key, value):
        return f"{key}={self._value_expr(key, value)}"

    def result_targets(self):
        return self.name

    def to_script(self):
        args = []
        for key, value in self.props.items():
            args.append(f"{key}={_literal(value)}")
        for key, value in self.arguments.items():
            args.append(self._argument_expr(key, value))
        return f"{self.result_targets()} = {_node_function_name(self.node)}({', '.join(args)})"


def _collect_node_props(node):
    props = {}
    parent_props = set()
    if getattr(node.bl_rna, "base", None) is not None:
        parent_props.update(prop.identifier for prop in node.bl_rna.base.bl_rna.properties)
    for prop in node.bl_rna.properties:
        if prop.identifier in parent_props or prop.identifier in _PROP_DENYLIST:
            continue
        if prop.is_readonly and prop.type not in {"ENUM", "BOOLEAN", "INT", "FLOAT", "STRING"}:
            continue
        try:
            value = getattr(node, prop.identifier)
        except Exception:
            continue
        if prop.type == "POINTER":
            continue
        if prop.type == "COLLECTION":
            continue
        props[_snake_case(prop.identifier)] = _convert_value(value)
    return props


def _collect_input_defaults(node):
    props = {}
    for socket in _enabled_sockets(node.inputs):
        if not hasattr(socket, "default_value") or socket.hide_value:
            continue
        props[_socket_name(socket)] = _convert_value(socket.default_value)
    return props


def _topological_sort(root):
    seen = set()
    order = []

    def visit(assignment):
        if assignment in seen:
            return
        seen.add(assignment)
        for dependency in assignment.flattened_arguments():
            visit(dependency)
        order.append(assignment)

    visit(root)
    return order


def _function_arguments(group_input):
    if group_input is None:
        return []
    arguments = []
    for output in _enabled_sockets(group_input.outputs):
        arguments.append(f"{_socket_name(output)}: {_socket_type_annotation(output)}")
    return arguments


def _return_entries(group_output, output_assignment):
    entries = []
    for socket in _enabled_sockets(group_output.node.inputs):
        key = _socket_name(socket)
        if key not in group_output.arguments:
            continue
        value = group_output.arguments[key]
        expr = group_output._argument_expr(key, value).split("=", 1)[1]
        entries.append(f'"{socket.name}": {expr}')
    if entries:
        return "{ " + ", ".join(entries) + " }"
    return output_assignment.name


def convert_node_tree(tree):
    """Return a best-effort Geometry Script draft for a Blender node tree."""
    assignments = []
    node_to_assignment = {}
    node_type_counts = {}
    incoming_links = {}

    for node in tree.nodes:
        prefix = _snake_case(node.type[:1] or "node")
        count = node_type_counts.get(prefix, 0) + 1
        node_type_counts[prefix] = count
        assignment = _Assignment(f"{prefix}{count}", node)
        assignment.props.update(_collect_node_props(node))
        assignments.append(assignment)
        node_to_assignment[node] = assignment

    for link in tree.links:
        incoming_links.setdefault(link.to_socket, []).append((node_to_assignment[link.from_node], _socket_name(link.from_socket)))

    for assignment in assignments:
        grouped_inputs = {}
        for socket in _enabled_sockets(assignment.node.inputs):
            key = _socket_name(socket)
            grouped_inputs.setdefault(key, []).append(socket)
        for key, sockets in grouped_inputs.items():
            values = []
            output_names = {}
            is_multi_input = any(socket.is_multi_input for socket in sockets)
            for socket in sockets:
                links = incoming_links.get(socket, [])
                if links:
                    for from_assignment, output_name in links:
                        output_names[len(values)] = output_name
                        values.append(from_assignment)
                elif hasattr(socket, "default_value") and not socket.hide_value:
                    values.append(_convert_value(socket.default_value))
            if not values:
                continue
            # Some Blender nodes expose the same conceptual control as both an
            # RNA enum property and an input socket. Geometry Script maps that
            # collapsed control to one Python keyword, so prefer the socket
            # value/link and drop the duplicate property keyword.
            assignment.props.pop(key, None)
            if len(values) == 1:
                assignment.arguments[key] = values[0]
                if output_names:
                    assignment.argument_outputs[key] = output_names[0]
            else:
                assignment.arguments[key] = values
                assignment.argument_outputs[key] = output_names
            assignment.argument_is_multi_input[key] = is_multi_input

    group_output = next((assignment for assignment in assignments if assignment.node.type == "GROUP_OUTPUT"), None)
    if group_output is None:
        raise ValueError(f"Node tree {tree.name!r} has no Group Output node")
    group_input = next((node for node in tree.nodes if node.type == "GROUP_INPUT"), None)

    sorted_assignments = [
        assignment
        for assignment in _topological_sort(group_output)
        if assignment.node.type not in {"GROUP_INPUT", "GROUP_OUTPUT"}
    ]
    body = "\n    ".join(assignment.to_script() for assignment in sorted_assignments)
    if not body:
        body = "pass"

    arguments = ", ".join(_function_arguments(group_input))
    function_name = _snake_case(tree.name)
    return_expr = _return_entries(group_output, sorted_assignments[-1] if sorted_assignments else group_output)

    return (
        "from geometry_script import *\n\n"
        f'@tree("{tree.name}")\n'
        f"def {function_name}({arguments}):\n"
        f"    {body}\n"
        f"    return {return_expr}\n"
    )


class ConvertTree(bpy.types.Operator):
    bl_idname = "geometry_script.convert_tree"
    bl_label = "Convert to Geometry Script"

    @classmethod
    def poll(cls, context):
        return bool(getattr(getattr(context, "space_data", None), "node_tree", None))

    def execute(self, context):
        script = convert_node_tree(context.space_data.node_tree)
        script_datablock = bpy.data.texts.new(context.space_data.node_tree.name + ".py")
        script_datablock.write(script)
        return {"FINISHED"}