diff --git a/nodes.py b/nodes.py index 9fb4f1e..73d8e76 100644 --- a/nodes.py +++ b/nodes.py @@ -4,7 +4,7 @@ import functools from collections import deque from copy import copy from math import degrees, radians -from typing import TypeVar, Callable, Optional, Iterable +from typing import TypeVar, Callable, Optional, Iterable, Tuple, List import bpy import nodeitems_utils @@ -13,6 +13,10 @@ from mathutils import Vector, Quaternion, Matrix, Euler T = TypeVar("T") +def _on_num_sockets_change(self: "SequenceNode", context: bpy.types.Context) -> None: + self.recreate(context) + + class RigNodesNodeTree(bpy.types.NodeTree): """Control EVERYTHING""" @@ -927,6 +931,137 @@ class Math(AbstractRigNodesNode): self.outputs["Result"].default_value = res +_enum_condition_data_types = [ + ("VECTOR", "Vector", ""), + ("ROTATION", "Rotation", ""), + ("FLOAT", "Float", ""), + ("INT", "Int", ""), +] + + +def _enum_dyn_conditional_value( + self: "ConditionalValueNode", context: bpy.types.Context +) -> List[Tuple[str, str, str]]: + if self.num_socks <= 0: + return [("NONE", "None", "")] + return [(f"{i}", f"Socket {i+1}", "") for i in range(0, self.num_socks)] + + +class ConditionalValueNode(AbstractRigNodesNode): + bl_idname = "ConditionalValueNode" + bl_label = "Conditional Value" + bl_icon = "EMPTY_ARROWS" + + dtype: bpy.props.EnumProperty( # type: ignore + name="Type", + items=_enum_condition_data_types, + default="VECTOR", + update=_on_num_sockets_change, + ) + num_socks: bpy.props.IntProperty( # type: ignore + default=2, + name="Sockets", + update=_on_num_sockets_change, + ) + active_sel: bpy.props.EnumProperty( # type: ignore + name="Active", + items=_enum_dyn_conditional_value, + ) + + def draw_buttons( + self, context: bpy.types.Context, layout: bpy.types.UILayout + ) -> None: + super().draw_buttons(context, layout) + layout.prop(self, "active_sel") + layout.prop(self, "dtype") + layout.prop(self, "num_socks") + + def init(self, context): + for index in range(0, self.num_socks): + match self.dtype: + case "VECTOR": + self.inputs.new("NodeSocketVector", f"Vector {index+1}") + case "ROTATION": + self.inputs.new("NodeSocketQuaternion", f"Rotation {index+1}") + case "FLOAT": + self.inputs.new("NodeSocketFloat", f"Float {index+1}") + case "INT": + self.inputs.new("NodeSocketInt", f"Int {index+1}") + + match self.dtype: + case "VECTOR": + self.outputs.new("NodeSocketVector", "Vector") + case "ROTATION": + self.outputs.new("NodeSocketQuaternion", "Rotation") + case "FLOAT": + self.outputs.new("NodeSocketFloat", "Float") + case "INT": + self.outputs.new("NodeSocketInt", "Int") + + def execute(self, depsgraph: bpy.types.Depsgraph) -> None: + match self.dtype: + case "VECTOR": + out = self._get_input_value(f"Vector {int(self.active_sel)+1}", Vector) + case "ROTATION": + out = self._get_input_value( + f"Rotation {int(self.active_sel)+1}", Quaternion + ) + case "FLOAT": + out = self._get_input_value(f"Float {int(self.active_sel)+1}", float) + case "INT": + out = self._get_input_value(f"Int {int(self.active_sel)+1}", int) + + self.outputs[0].default_value = out + + +class ConditionalExecutionNode(AbstractRigNodesNode): + bl_idname = "ConditionalExecutionNode" + bl_label = "Conditional Execution" + bl_icon = "EMPTY_ARROWS" + + num_socks: bpy.props.IntProperty( # type: ignore + default=2, + name="Sockets", + update=_on_num_sockets_change, + ) + active_sel: bpy.props.EnumProperty( # type: ignore + name="Active", + items=_enum_dyn_conditional_value, + ) + + def draw_buttons( + self, context: bpy.types.Context, layout: bpy.types.UILayout + ) -> None: + super().draw_buttons(context, layout) + layout.prop(self, "active_sel") + layout.prop(self, "num_socks") + + def init(self, context: bpy.types.Context) -> None: + self.inputs.new(NodeSocketExecute.bl_idname, NodeSocketExecute.bl_label) + for index in range(0, self.num_socks): + self.outputs.new( + NodeSocketExecute.bl_idname, f"{NodeSocketExecute.bl_label} {index + 1}" + ) + + def exec_order_successors(self) -> Iterable["AbstractRigNodesNode"]: + """Generator, yields the nodes that should be executed after this one.""" + + # For output execution order, only consider actual 'Execute' sockets. + def follow_socket(socket: bpy.types.NodeSocket) -> bool: + return isinstance(socket, NodeSocketExecute) or isinstance( + socket.node, AbstractAlwaysExecuteNode + ) + + if self.active_sel is not "NONE": + return self._connected_nodes( + [self.outputs[int(self.active_sel)]], "to_socket", follow_socket + ) + return [] + + def execute(self, depsgraph: bpy.types.Depsgraph) -> None: + pass + + class SetCursorNode(AbstractAlwaysExecuteNode): """Sets the location and/or rotation of the 3D cursor""" @@ -1232,10 +1367,6 @@ class ClampNode(AbstractRigNodesNode): self.outputs["Result"].default_value = clamped -def _on_num_sockets_change(self: "SequenceNode", context: bpy.types.Context) -> None: - self.recreate(context) - - class SequenceNode(AbstractRigNodesNode): """Multiple 'Execute' node sockets.""" @@ -1307,6 +1438,8 @@ node_categories = [ "Flow", items=[ nodeitems_utils.NodeItem("SequenceNode"), + nodeitems_utils.NodeItem("ConditionalValueNode"), + nodeitems_utils.NodeItem("ConditionalExecutionNode"), ], ), RigNodesNodeCategory( @@ -1371,6 +1504,8 @@ classes = ( TwoBoneIKNode, SetCursorNode, SequenceNode, + ConditionalValueNode, + ConditionalExecutionNode, # Math Nodes RotateTowards, AngleFromVectors,