Merge b78ec37f0f into 96a1b293cc

__init__.py

@@ -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:
@@ -98,7 +103,10 @@ def register():
     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 +119,11 @@ 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:

operators/convert_tree.py

@@ -0,0 +1,148 @@
+import bpy
+import mathutils
+import collections.abc
+
+class _Assignment:
+    name: str
+    node: bpy.types.Node
+    props: dict
+    arguments: dict
+    argument_dot_access: dict
+
+    def __init__(self, name, node, props):
+        self.name = name
+        self.node = node
+        self.props = props
+        self.arguments = {}
+        self.argument_dot_access = {}
+    
+    def _flattened_arguments(self):
+        for a in self.arguments.values():
+            if isinstance(a, _Assignment):
+                yield a
+            elif isinstance(a, list):
+                yield from a
+    def flattened_arguments(self):
+        return [*self._flattened_arguments()]
+
+    def convert_argument(self, k, v, delimiter='=', index=None):
+        if isinstance(v, list):
+            v = ', '.join([self.convert_argument(k, sv, delimiter="", index=i).removeprefix(k) for i, sv in enumerate(v)])
+            return f"{k}{delimiter}[{v}]"
+        if not isinstance(v, _Assignment):
+            if isinstance(v, str):
+                v = f'"{v}"'
+            else:
+                v = str(v)
+            return f"{k}{delimiter}{v}"
+        if v.node.type == 'GROUP_INPUT':
+            return f"{k}{delimiter}{self.argument_dot_access[k] if index is None else self.argument_dot_access[k][index]}"
+        else:
+            return f"{k}{delimiter}{v.name}{'.' + (self.argument_dot_access[k] if index is None else self.argument_dot_access[k][index]) if len(list(o for o in v.node.outputs if o.enabled)) > 1 else ''}"
+
+    def to_script(self):
+        snake_case_name = self.node.bl_rna.name.lower().replace(' ', '_')
+        args = ', '.join([
+            self.convert_argument(k, v)
+
+            for k, v in (list(self.props.items()) + list(self.arguments.items()))
+        ])
+        return f"{self.name} = {snake_case_name}({args})"
+
+class ConvertTree(bpy.types.Operator):
+    bl_idname = "geometry_script.convert_tree"
+    bl_label = "Convert to Geometry Script"
+
+    def execute(self, context):
+        tree = context.space_data.node_tree
+
+        assignments = []
+
+        def convert_default_value(v):
+            if isinstance(v, mathutils.Vector):
+                return (v[0], v[1], v[2])
+            elif isinstance(v, mathutils.Euler):
+                return (v[0], v[1], v[2])
+            elif isinstance(v, collections.abc.Iterable):
+                return tuple(v)
+            else:
+                return v
+
+        node_type_counter = {}
+        for node in tree.nodes:
+            count = node_type_counter.get(node.type[0], 0)
+            props = {i.name.lower().replace(' ', '_'): convert_default_value(i.default_value) for i in node.inputs if i.enabled and hasattr(i, 'default_value') and not i.hide_value}
+            props.update({k: getattr(node, k) for k in set(p.identifier for p in node.bl_rna.properties) - set(p.identifier for p in node.bl_rna.base.bl_rna.properties)})
+            assignments.append(_Assignment(f"{node.type[0].lower()}{count + 1}", node, props))
+            node_type_counter[node.type[0]] = count + 1
+
+        sorted_assignments = assignments.copy()
+        for link in tree.links:
+            to_node = next(a for a in assignments if a.node == link.to_node)
+            from_node = next(a for a in assignments if a.node == link.from_node)
+            argument_name = link.to_socket.name.lower().replace(' ', '_')
+            output_name = link.from_socket.name.lower().replace(' ', '_')
+            if argument_name in to_node.props:
+                del to_node.props[argument_name]
+            if argument_name in to_node.arguments:
+                if isinstance(to_node.arguments[argument_name], list):
+                    index = len(to_node.arguments[argument_name])
+                    to_node.arguments[argument_name].append(from_node)
+                    to_node.argument_dot_access[argument_name][index] = output_name
+                else:
+                    to_node.arguments[argument_name] = [to_node.arguments[argument_name], from_node]
+                    to_node.argument_dot_access[argument_name] = {
+                        0: to_node.argument_dot_access[argument_name],
+                        1: output_name
+                    }
+            else:
+                to_node.arguments[argument_name] = from_node
+                to_node.argument_dot_access[argument_name] = output_name
+
+        def topological_sort(root):
+            seen = set()
+            stack = []
+            order = []
+            q = [root]
+            while q:
+                v = q.pop()
+                if v not in seen:
+                    seen.add(v)
+                    q.extend(v.flattened_arguments())
+
+                    while stack and v not in stack[-1].flattened_arguments():
+                        order.append(stack.pop())
+                    stack.append(v)
+
+            return order[::-1] + stack[::-1]
+        
+        group_output = next(a for a in assignments if a.node.type == 'GROUP_OUTPUT')
+        sorted_assignments = [a for a in topological_sort(group_output) if a.node.type != 'GROUP_OUTPUT' and a.node.type != 'GROUP_INPUT']
+        
+        body = '\n    '.join([a.to_script() for a in sorted_assignments])
+
+        group_input = next((n for n in tree.nodes if n.type == 'GROUP_INPUT'), None)
+        tree_arguments = []
+        if group_input is not None:
+            for output in group_input.outputs:
+                if output.type == 'CUSTOM':
+                    continue
+                output_name = output.name.lower().replace(' ', '_')
+                output_type = output.bl_rna.identifier.replace('NodeSocket', '')
+                tree_arguments.append(f"{output_name}: {output_type}")
+        tree_arguments = ', '.join(tree_arguments)
+
+        tree_returns = []
+        for k, v in group_output.arguments.items():
+            tree_returns.append(group_output.convert_argument(f'"{k}"', v, ': '))
+        tree_returns = ', '.join(tree_returns)
+
+        script = f"""from geometry_script import *
+
+@tree("{tree.name}")
+def {tree.name.lower().replace(' ', '_')}({tree_arguments}):
+    {body}
+    return {{ {tree_returns} }}"""
+        script_datablock = bpy.data.texts.new(tree.name + '.py')
+        script_datablock.write(script)
+        return {'FINISHED'}