Port node tree conversion operator
This commit is contained in:
Meta Vens
parent
55085d0a4b
commit
277d8b0128
3 changed files with 326 additions and 0 deletions
|
|
@ -22,6 +22,7 @@ import webbrowser
|
|||
from .api.tree import *
|
||||
from .preferences import GeometryScriptPreferences
|
||||
from .absolute_path import absolute_path
|
||||
from .operators.convert_tree import ConvertTree
|
||||
|
||||
bl_info = {
|
||||
"name" : "Geometry Script",
|
||||
|
|
@ -74,6 +75,10 @@ def templates_menu_draw(self, context):
|
|||
def editor_header_draw(self, context):
|
||||
self.layout.menu(GeometryScriptMenu.bl_idname)
|
||||
|
||||
def node_header_draw(self, context):
|
||||
if context.space_data.tree_type == 'GeometryNodeTree':
|
||||
self.layout.operator(ConvertTree.bl_idname)
|
||||
|
||||
def auto_resolve():
|
||||
if bpy.context.scene.geometry_script_settings.auto_resolve:
|
||||
try:
|
||||
|
|
@ -97,8 +102,10 @@ def register():
|
|||
bpy.utils.register_class(GeometryScriptPreferences)
|
||||
bpy.utils.register_class(OpenDocumentation)
|
||||
bpy.utils.register_class(GeometryScriptMenu)
|
||||
bpy.utils.register_class(ConvertTree)
|
||||
|
||||
bpy.types.TEXT_HT_header.append(editor_header_draw)
|
||||
bpy.types.NODE_HT_header.append(node_header_draw)
|
||||
|
||||
bpy.types.Scene.geometry_script_settings = bpy.props.PointerProperty(type=GeometryScriptSettings)
|
||||
|
||||
|
|
@ -111,7 +118,9 @@ def unregister():
|
|||
bpy.utils.unregister_class(GeometryScriptPreferences)
|
||||
bpy.utils.unregister_class(OpenDocumentation)
|
||||
bpy.utils.unregister_class(GeometryScriptMenu)
|
||||
bpy.utils.unregister_class(ConvertTree)
|
||||
bpy.types.TEXT_HT_header.remove(editor_header_draw)
|
||||
bpy.types.NODE_HT_header.remove(node_header_draw)
|
||||
try:
|
||||
bpy.app.timers.unregister(auto_resolve)
|
||||
except:
|
||||
|
|
|
|||
1
operators/__init__.py
Normal file
1
operators/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
"""Geometry Script Blender operators."""
|
||||
316
operators/convert_tree.py
Normal file
316
operators/convert_tree.py
Normal file
|
|
@ -0,0 +1,316 @@
|
|||
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
|
||||
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"}
|
||||
Reference in a new issue