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,