Refactor simulation API to use simulation_zone decorator

api/static/simulation.py

@@ -1,25 +1,46 @@
+import bpy
 import inspect
 import typing
 
-class SimulationInput:
-    class DeltaTime: pass
-    class ElapsedTime: pass
-
-def simulation(block: typing.Callable[typing.Any, 'Geometry']):
+def simulation_zone(block: typing.Callable):
     """
     Create a simulation input/output block.
 
-    > Only available in the `geometry-node-simulation` branch of Blender 3.5.
+    > Only available in Blender 3.6+.
     """
-    def wrapped(geometry: 'Geometry', *args, **kwargs):
-        from geometry_script import simulation_input, simulation_output
-        simulation_in = simulation_input(geometry=geometry)
+    def wrapped(*args, **kwargs):
+        from geometry_script.api.node_mapper import OutputsList, set_or_create_link
+        from geometry_script.api.state import State
+        from geometry_script.api.types import Type, socket_class_to_data_type
+
         signature = inspect.signature(block)
-        for key, value in signature.parameters.items():
-            match value.annotation:
-                case SimulationInput.DeltaTime:
-                    kwargs[key] = simulation_in.delta_time
-                case SimulationInput.ElapsedTime:
-                    kwargs[key] = simulation_in.elapsed_time
-        return simulation_output(geometry=block(simulation_in.geometry, *args, **kwargs)).geometry
+        
+        # setup zone
+        simulation_in = State.current_node_tree.nodes.new(bpy.types.GeometryNodeSimulationInput.__name__)
+        simulation_out = State.current_node_tree.nodes.new(bpy.types.GeometryNodeSimulationOutput.__name__)
+        simulation_in.pair_with_output(simulation_out)
+
+        # clear state items
+        for item in simulation_out.state_items:
+            simulation_out.state_items.remove(item)
+
+        # create state items from block signature
+        state_items = {}
+        for param in [*signature.parameters.values()][1:]:
+            state_items[param.name] = (param.annotation, param.default, None, None)
+        for i, arg in enumerate(state_items.items()):
+            simulation_out.state_items.new(socket_class_to_data_type(arg[1][0].socket_type), arg[0].replace('_', ' ').title())
+            set_or_create_link(kwargs[arg[0]] if arg[0] in kwargs else args[i], simulation_in.inputs[i])
+        
+        step = block(*[Type(o) for o in simulation_in.outputs[:-1]])
+
+        if isinstance(step, Type):
+            step = (step,)
+        for i, result in enumerate(step):
+            State.current_node_tree.links.new(result._socket, simulation_out.inputs[i])
+
+        if len(simulation_out.outputs[:-1]) == 1:
+            return Type(simulation_out.outputs[0])
+        else:
+            return OutputsList({o.name.lower().replace(' ', '_'): Type(o) for o in simulation_out.outputs[:-1]})
     return wrapped