126 lines
No EOL
3.9 KiB
Python
126 lines
No EOL
3.9 KiB
Python
from __future__ import annotations
|
|
import typing
|
|
|
|
from typing import Any
|
|
from bpy.types import NodeSocket, Node
|
|
|
|
class NodeException(Exception): ...
|
|
|
|
_EvalCache = dict[tuple[int, str], Any]
|
|
|
|
class AutoBlIdName:
|
|
def __init_subclass__(cls, **kwargs):
|
|
super().__init_subclass__(**kwargs)
|
|
if "bl_idname" not in cls.__dict__:
|
|
cls.bl_idname = cls.__name__
|
|
|
|
class SliceNodeSocket(NodeSocket, AutoBlIdName):
|
|
default: Any = None
|
|
color = (1., 1., 1., 1.)
|
|
|
|
def draw_color(self, context, node) -> tuple[float, float, float, float]: return self.color # pyright: ignore[reportIncompatibleMethodOverride]
|
|
|
|
class SliceNode(Node, AutoBlIdName):
|
|
def evaluate_output(self, socket_name: str, **_kwargs) -> Any: ...
|
|
|
|
def add_output(self, bl_idname: str, name: str | None = None):
|
|
from ..registry import get_label_by_name
|
|
return self.outputs.new(bl_idname, name or get_label_by_name(bl_idname))
|
|
|
|
def add_input(self, bl_idname: str, name: str | None = None, multi: bool = False):
|
|
from ..registry import get_label_by_name
|
|
sock = self.inputs.new(bl_idname, name or get_label_by_name(bl_idname))
|
|
if multi:
|
|
sock.link_limit = 0
|
|
sock.display_shape = 'LIST'
|
|
return sock
|
|
|
|
## Evaluation helpers
|
|
|
|
def _follow_reroute(socket: NodeSocket) -> NodeSocket:
|
|
current: NodeSocket = socket
|
|
seen: set[int] = set()
|
|
|
|
while getattr(current.node, "bl_idname", "") == "NodeReroute":
|
|
assert current.node
|
|
ptr = current.node.as_pointer()
|
|
if ptr in seen:
|
|
break
|
|
seen.add(ptr)
|
|
|
|
reroute_input = current.node.inputs[0] if current.node.inputs else None
|
|
if not reroute_input or not reroute_input.is_linked:
|
|
raise NodeException("Loose reroute node found: ensure all reroute nodes are correctly linked")
|
|
|
|
current = reroute_input.links[0].from_socket
|
|
|
|
return current
|
|
|
|
def resolve_input(
|
|
socket: SliceNodeSocket | None,
|
|
*,
|
|
cache: _EvalCache | None = None,
|
|
visited: set[tuple[int, str]] | None = None,
|
|
):
|
|
if socket is None:
|
|
raise Exception(f"The queried socket is non-existant")
|
|
|
|
assert socket.node
|
|
|
|
if not socket.is_linked:
|
|
if v := getattr(socket, "value", None): return v
|
|
else: return socket.default
|
|
|
|
assert socket.is_linked
|
|
assert socket.links
|
|
from_socket: NodeSocket = _follow_reroute(socket.links[0].from_socket)
|
|
|
|
return _resolve_output(
|
|
from_socket,
|
|
default=None,
|
|
cache=cache,
|
|
visited=visited,
|
|
)
|
|
|
|
def _resolve_output(
|
|
socket: NodeSocket,
|
|
*,
|
|
default: Any = None,
|
|
cache: _EvalCache | None = None,
|
|
visited: set[tuple[int, str]] | None = None,
|
|
):
|
|
|
|
cache = {} if cache is None else cache
|
|
visited = set() if visited is None else visited
|
|
|
|
assert socket.node
|
|
key = (socket.node.as_pointer(), socket.node.name)
|
|
if key in cache: return cache[key]
|
|
if key in visited: return default # cycle guard
|
|
|
|
visited.add(key)
|
|
try:
|
|
if not isinstance(socket.node, SliceNode):
|
|
return default
|
|
|
|
if getattr(socket.node, "mute", False):
|
|
# Attempt bypass: find a matching input socket by name or type
|
|
input_to_bypass = socket.node.inputs.get(socket.name)
|
|
if not input_to_bypass:
|
|
for inp in socket.node.inputs:
|
|
if inp.bl_idname == socket.bl_idname:
|
|
input_to_bypass = inp
|
|
break
|
|
|
|
if input_to_bypass:
|
|
value = resolve_input(typing.cast(Any, input_to_bypass), cache=cache, visited=visited)
|
|
cache[key] = value
|
|
return value
|
|
return default
|
|
|
|
assert socket.node.evaluate_output
|
|
value = socket.node.evaluate_output(socket.name, cache=cache, visited=visited)
|
|
cache[key] = value
|
|
return value
|
|
finally:
|
|
visited.remove(key) |