Loop Nodes #4

Open
Denys Hsu wants to merge 5 commits from cgtinker/powership:loop into main

When changing the target branch, be careful to rebase the branch in your fork to match. See documentation.
4 changed files with 339 additions and 13 deletions
Showing only changes of commit 9eccd9d8aa - Show all commits

1
gui.py
View File

@ -54,6 +54,7 @@ def draw_nodetree_header_menu(self, context):
layout = self.layout
row = layout.row(align=False)
row.operator("rignodes.execute_tree", icon="PLAY")
row.operator("rignodes.prepare_tree", icon="NODETREE")
row.prop(tree, "autorun")
row.prop(context.scene, "rignodes_mode")
row.operator("rignodes.rebuild_node")

326
nodes.py
View File

@ -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, Set
import bpy
import nodeitems_utils
@ -65,21 +65,78 @@ class RigNodesNodeTree(bpy.types.NodeTree):
return node
return None
def _prepare_nodes(self) -> None:
def _reset_nodes(self) -> None:
for node in self.nodes:
node.reset_run()
def _prepare_nodes(self) -> None:
self._reset_nodes()
# Find all loops in the tree.
output_nodes = list()
for node in self.nodes:
if isinstance(node, LoopOutputNode):
output_nodes.append(node)
cgtinker marked this conversation as resolved

This can be written as:

output_nodes = [node for node in self.nodes
                if isinstance(node, LoopOutputNode)]

And since it's only being iterated over once, there even no need to construct a list, so a generator expression does the trick as well:

output_nodes = (node for node in self.nodes
                if isinstance(node, LoopOutputNode))
This can be written as: ```py output_nodes = [node for node in self.nodes if isinstance(node, LoopOutputNode)] ``` And since it's only being iterated over once, there even no need to construct a list, so a generator expression does the trick as well: ```py output_nodes = (node for node in self.nodes if isinstance(node, LoopOutputNode)) ```
# Set iteration count foreach (nested) loop.
for output_node in output_nodes:
input_node = output_node._get_connection_node()
nodes = list(output_node._get_nodes_between(input_node))
for node in nodes:
cgtinker marked this conversation as resolved

There is no need to construct a list here, just loop over the generator returned by output_node._get_loop_nodes()

There is no need to construct a list here, just loop over the generator returned by `output_node._get_loop_nodes()`
if node == output_node:
continue
node.iterations *= output_node.num_socks
def _run_loop(
self,
node: "LoopInputNode",
depsgraph: bpy.types.Depsgraph,
to_visit: deque["AbstractRigNodesNode"],
depth: int = 1,
) -> None:
"""Run one iteration of a loop.
:return: Node(s) to visit next."""
cgtinker marked this conversation as resolved

Formatting: when writing multi-line docstrings, the closing """ should be on a line of its own (see PEP 257)

Formatting: when writing multi-line docstrings, the closing `"""` should be on a line of its own (see [PEP 257](https://peps.python.org/pep-0257/#multi-line-docstrings))
Review

Sorry, still stumbling.

Sorry, still stumbling.
indent = '\t' * depth
print(f"{indent}\033[38;5;214mRunning {node}\033[0m")
input_node = node
output_node = node._get_connection_node()
while to_visit:
node = to_visit.popleft()
nested_loop: Optional[deque[AbstractRigNodesNode]] = None
if isinstance(node, LoopInputNode) and node != input_node:
# Run nested loop.
nested_output_node = node._get_connection_node()
nested_loop = deque(nested_output_node._get_nodes_between(node))
nested_loop.append(nested_output_node)
self._run_loop(node, depsgraph, nested_loop, depth+1)
# continue
print(f"{indent}\t-{node}")
node.run(depsgraph)
# if node.iterations > 0:
# # Add output node for future iterations.
# to_visit.append(node)
if node == output_node:
break
def _run_from_node(
self, depsgraph: bpy.types.Depsgraph, start_node: "AbstractRigNodesNode"
) -> None:
# Execute the ForwardSolveNode and traverse its output connections
to_visit = deque([start_node])
to_visit: deque[AbstractRigNodesNode] = deque([start_node])
visited = set()
print(f"\033[38;5;214mRunning tree {self.name}\033[0m")
while to_visit:
node: AbstractRigNodesNode = to_visit[0]
print(f" - {node}")
print(f"\t- {node}")
nodes_before = list(node.exec_order_prerequisites())
if nodes_before:
@ -88,6 +145,11 @@ class RigNodesNodeTree(bpy.types.NodeTree):
to_visit.extendleft(reversed(nodes_before))
continue
if isinstance(node, LoopInputNode):
# Run one iteration if the found loop (recursive if nested).
self._run_loop(node, depsgraph, to_visit)
continue
# Everything that this node depends on has run, so time to run it
# itself. It can be taken off the queue.
to_visit.popleft()
@ -164,6 +226,7 @@ class AbstractRigNodesNode(bpy.types.Node):
description="Error message; if empty, everything is fine",
update=_on_node_error_updated,
)
iterations: bpy.props.IntProperty(name="Iterations", default=1) # type: ignore
def init(self, context: bpy.types.Context) -> None:
pass
@ -179,6 +242,9 @@ class AbstractRigNodesNode(bpy.types.Node):
else:
layout.label(text="did not run", icon="PAUSE")
if context.preferences.view.show_developer_ui:
layout.prop(self, "iterations")
if self.error:
layout.label(text=self.error, icon="ERROR")
@ -199,10 +265,29 @@ class AbstractRigNodesNode(bpy.types.Node):
sock.hide_value = True
return sock
@property
def _inputs(self):
return self.inputs
@property
def _outputs(self):
return self.outputs
def _get_nodes_between(
self, to_node: "AbstractRigNodesNode"
) -> Iterable["AbstractRigNodesNode"]:
# TODO: the naming is bad here I think. Check that it's actually from -> to or to -> from
nodes = list(self._connected_nodes(self._inputs, "from_socket"))
if not to_node in nodes and len(nodes) > 0:
from_node = nodes[-1]
yield from from_node._get_nodes_between(to_node)
yield from nodes
def exec_order_prerequisites(self) -> Iterable["AbstractRigNodesNode"]:
"""Generator, yields the nodes that should be executed before this one."""
# For input execution order, consider all connected nodes regardless of socket types.
return self._connected_nodes(self.inputs, "from_socket")
return self._connected_nodes(self._inputs, "from_socket")
def exec_order_successors(self) -> Iterable["AbstractRigNodesNode"]:
"""Generator, yields the nodes that should be executed after this one."""
@ -213,7 +298,7 @@ class AbstractRigNodesNode(bpy.types.Node):
socket.node, AbstractAlwaysExecuteNode
)
return self._connected_nodes(self.outputs, "to_socket", follow_socket)
return self._connected_nodes(self._outputs, "to_socket", follow_socket)
def _connected_nodes(
self,
@ -259,6 +344,8 @@ class AbstractRigNodesNode(bpy.types.Node):
self.error = str(ex)
raise
finally:
self.iterations -= 1
if self.iterations == 0:
self.has_run = True
def __str__(self) -> str:
@ -267,6 +354,7 @@ class AbstractRigNodesNode(bpy.types.Node):
def reset_run(self) -> None:
self.has_run = False
self.error = ""
self.iterations = 1
def _get_input_socket_value(self, input_socket: bpy.types.NodeSocket) -> object:
for link in input_socket.links:
@ -415,6 +503,222 @@ class AbstractRigNodesNode(bpy.types.Node):
tree.links.new(from_socket, to_socket)
def _on_update_recreate_node(self: "SequenceNode", context: bpy.types.Context) -> None:
self.recreate(context)
_enum_loop_dtype = (
("OBJECT", "Object", ""),
("VECTOR", "Vector", ""),
("QUATERNION", "Rotation", ""),
("VALUE", "Value", ""),
)
class LoopInputNode(AbstractRigNodesNode):
bl_idname = "LoopInputNode"
bl_label = "Loop Input Node"
bl_icon = "LOOP_FORWARDS"
dtype: bpy.props.EnumProperty( # type: ignore
name="Type",
items=_enum_loop_dtype,
default="OBJECT",
update=_on_update_recreate_node,
)
num_socks: bpy.props.IntProperty( # type: ignore
name="Number of Sockets",
default=2,
update=_on_update_recreate_node,
)
index: bpy.props.IntProperty(name="Index", default=0) # type: ignore
def draw_buttons(
self, context: bpy.types.Context, layout: bpy.types.UILayout
) -> None:
super().draw_buttons(context, layout)
layout.prop(self, "dtype")
layout.prop(self, "num_socks")
layout.prop(self, "index")
# 67649: The socket_value_update fn does not work when changing values in a node tree.
# Therefore using update for prototying.
def update(self):
if "Loop Connection" in self.outputs:
setattr(self.outputs["Loop Connection"], "default_value", self.num_socks)
def _get_connection_node(self) -> "LoopOutputNode":
# The 'Loop Connection' should be connected to a 'Loop Input Node',
# otherwise the loop won't be executed.
for socket in self.outputs[:1]:
for link in socket.links[:1]:
connected_socket = getattr(link, "to_socket")
if not hasattr(connected_socket, "node"):
raise RuntimeError(
"Loop Output Node is required to execute a loop."
)
if isinstance(connected_socket.node, LoopOutputNode):
return connected_socket.node
raise RuntimeError("Loop Output Node is required to execute a loop.")
def reset_run(self):
super().reset_run()
self.index = 0
def init(self, context: bpy.types.Context) -> None:
for index in range(self.num_socks):
match self.dtype:
case "OBJECT":
self.inputs.new("NodeSocketObject", f"Object {index+1}")
case "VECTOR":
self.inputs.new("NodeSocketVector", f"Vector {index+1}")
case "QUATERNION":
self.inputs.new("NodeSocketQuaternion", f"Rotation {index+1}")
case "VALUE":
self.inputs.new("NodeSocketFloat", f"Value {index+1}")
self.outputs.new("NodeSocketInt", "Loop Connection")
match self.dtype:
case "OBJECT":
self.outputs.new("NodeSocketObject", f"Object")
case "VECTOR":
self.outputs.new("NodeSocketVector", f"Vector")
case "QUATERNION":
self.outputs.new("NodeSocketQuaternion", f"Rotation")
case "VALUE":
self.outputs.new("NodeSocketFloat", f"Value")
def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
if self.index == self.num_socks:
self.index = 0
self.index += 1
index = self.index
match self.dtype:
case "OBJECT":
obj = self._get_input_value(f"Object {index}", bpy.types.Object)
self.outputs["Object"].default_value = obj
case "VECTOR":
value = self._get_input_value(f"Vector {index}", Vector)
self.outputs["Vector"].default_value = value
case "QUATERNION":
value = self._get_input_value(f"Quaternion {index}", Quaternion)
self.outputs["Rotation"].default_value = value
case "VALUE":
value = self._get_input_value(f"Value {index}", float)
self.outputs["Value"].default_value = value
class LoopOutputNode(AbstractRigNodesNode):
bl_idname = "LoopOutputNode"
bl_label = "Loop Output Node"
bl_icon = "LOOP_BACK"
dtype: bpy.props.EnumProperty( # type: ignore
name="Type",
items=_enum_loop_dtype,
default="OBJECT",
update=_on_update_recreate_node,
)
num_socks: bpy.props.IntProperty( # type: ignore
name="Number of Sockets",
default=2,
update=_on_update_recreate_node,
)
index: bpy.props.IntProperty(name="Index", default=0) # type: ignore
def draw_buttons(
self, context: bpy.types.Context, layout: bpy.types.UILayout
) -> None:
super().draw_buttons(context, layout)
layout.prop(self, "dtype")
layout.prop(self, "index")
def _get_connection_node(self) -> "LoopInputNode":
# The 'Loop Connection' should be connected to a 'Loop Input Node',
# otherwise the loop won't be executed.
for socket in self.inputs[:1]:
for link in socket.links[:1]:
connected_socket = getattr(link, "from_socket")
if not hasattr(connected_socket, "node"):
raise RuntimeError("Loop Input Node is required to execute a loop.")
if isinstance(connected_socket.node, LoopInputNode):
return connected_socket.node
raise RuntimeError("Loop Input Node is required to execute a loop.")
@AbstractRigNodesNode._inputs.getter # type: ignore
def _inputs(self):
return self.inputs[1:]
def init(self, context: bpy.types.Context) -> None:
self.inputs.new("NodeSocketInt", "Loop Connection")
match self.dtype:
case "OBJECT":
self.inputs.new("NodeSocketObject", f"Object")
case "VECTOR":
self.inputs.new("NodeSocketVector", f"Vector")
case "QUATERNION":
self.inputs.new("NodeSocketQuaternion", f"Rotation")
case "VALUE":
self.inputs.new("NodeSocketFloat", f"Value")
for index in range(self.num_socks):
match self.dtype:
case "OBJECT":
self.outputs.new("NodeSocketObject", f"Object {index+1}")
case "VECTOR":
self.outputs.new("NodeSocketVector", f"Vector {index+1}")
case "QUATERNION":
self.outputs.new("NodeSocketQuaternion", f"Rotation {index+1}")
case "VALUE":
self.outputs.new("NodeSocketFloat", f"Value {index+1}")
def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
if self.index == self.num_socks:
self.index = 0
self.index += 1
index = self.index
match self.dtype:
case "OBJECT":
obj = self._get_input_value(f"Object", bpy.types.Object)
self.outputs[f"Object {index}"].default_value = obj
case "VECTOR":
obj = self._get_input_value(f"Vector", Vector)
self.outputs[f"Vector {index}"].default_value = obj
case "QUATERNION":
obj = self._get_input_value(f"Rotation", Quaternion)
self.outputs[f"Rotation {index}"].default_value = obj
case "VALUE":
obj = self._get_input_value(f"Value", float)
self.outputs[f"Value {index}"].default_value = obj
def reset_run(self) -> None:
super().reset_run()
self.index = 0
self.iterations = self.num_socks
# 67649: The socket_value_update fn does not work when changing values in a node tree.
# Therefore using update for prototying the loop function.
# This allows to dynamically change the sockets based loop input socket amount.
def update(self):
if "Loop Connection" not in self.inputs:
return
val = self._get_optional_input_value("Loop Connection", int)
if val != self.num_socks and isinstance(val, int):
self.num_socks = val
class AbstractRigNodesEventNode(AbstractRigNodesNode):
"""Node that can only exist once in a tree."""
@ -1232,10 +1536,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."""
@ -1246,7 +1546,7 @@ class SequenceNode(AbstractRigNodesNode):
num_socks: bpy.props.IntProperty( # type: ignore
name="Number of Sockets",
default=2,
update=_on_num_sockets_change,
update=_on_update_recreate_node,
)
def draw_buttons(
@ -1307,6 +1607,8 @@ node_categories = [
"Flow",
items=[
nodeitems_utils.NodeItem("SequenceNode"),
nodeitems_utils.NodeItem("LoopInputNode"),
nodeitems_utils.NodeItem("LoopOutputNode"),
],
),
RigNodesNodeCategory(
@ -1371,6 +1673,8 @@ classes = (
TwoBoneIKNode,
SetCursorNode,
SequenceNode,
LoopInputNode,
LoopOutputNode,
# Math Nodes
RotateTowards,
AngleFromVectors,

View File

@ -56,6 +56,27 @@ class RigNodes_OT_snap_to_bone(bpy.types.Operator):
snap_to_bone(context.object, context.view_layer.depsgraph)
return {"FINISHED"}
class RigNodes_OT_prepare_tree(bpy.types.Operator):
"""Prepare node tree manually."""
classes = (RigNodes_OT_snap_to_bone,)
bl_idname = "rignodes.prepare_tree"
bl_label = "Prepare Tree"
bl_options = {"REGISTER", "UNDO"}
@classmethod
def poll(cls, context: bpy.types.Context) -> bool:
return bool(context.object) and context.preferences.view.show_developer_ui
def execute(self, context: bpy.types.Context) -> set[str]:
for tree in bpy.data.node_groups:
if tree.bl_idname != "RigNodesNodeTree":
continue
if tree.is_running:
continue
tree._prepare_nodes()
return {"FINISHED"}
classes = (RigNodes_OT_snap_to_bone, RigNodes_OT_prepare_tree,)
register, unregister = bpy.utils.register_classes_factory(classes)

BIN
rignode-loops.blend Normal file

Binary file not shown.