228 lines
7.8 KiB
Python
228 lines
7.8 KiB
Python
from __future__ import annotations
|
|
from dataclasses import field
|
|
import typing
|
|
import bpy
|
|
from contextlib import contextmanager
|
|
|
|
from typing import Any
|
|
from bpy.types import NodeSocket, Node
|
|
|
|
class NodeException(Exception): ...
|
|
|
|
_EvalCache = dict[tuple[int, str, 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
|
|
bl_color = (1.0, 1.0, 1.0, 1.0)
|
|
|
|
def draw_color(self, context, node) -> tuple[float, float, float, float]: # pyright: ignore[reportIncompatibleMethodOverride]
|
|
return self.bl_color
|
|
|
|
from ...registry import get_label_by_name
|
|
class SliceNode(Node, AutoBlIdName):
|
|
def update(self) -> None:
|
|
self._update_multi_inputs()
|
|
|
|
def evaluate_output(self, socket_name: str, **_kwargs) -> Any: ...
|
|
|
|
def resolve(self, fieldname) -> Any:
|
|
cache = getattr(self, "_current_cache", None)
|
|
visited = getattr(self, "_current_visited", None)
|
|
return resolve_input(self.inputs.get(fieldname), cache=cache, visited=visited)
|
|
|
|
def add_output(self, bl_idname: str, name: str = "", identifier: str = ""):
|
|
eval_name = name or get_label_by_name(bl_idname)
|
|
return self.outputs.new(bl_idname, eval_name, identifier=identifier or eval_name)
|
|
|
|
def add_input(self, bl_idname: str, name: str = "", identifier: str = ""):
|
|
eval_name = name or get_label_by_name(bl_idname)
|
|
return self.inputs.new(bl_idname, eval_name, identifier=identifier or eval_name)
|
|
|
|
def add_multi_input(self, bl_idname: str, name: str = "") -> None:
|
|
if not self.get('multi_sockets'): self['multi_sockets'] = {}
|
|
if not self['multi_sockets'].get(name): self['multi_sockets'][name] = {'bl_idname': bl_idname, 'name': name or get_label_by_name(bl_idname)}
|
|
self._update_multi_inputs()
|
|
|
|
def resolve_multi_input(self, fieldname) -> Any:
|
|
cache = getattr(self, "_current_cache", None)
|
|
visited = getattr(self, "_current_visited", None)
|
|
inputs = [i for i in list(self.inputs) if i.name == fieldname]
|
|
return [resolve_input(r, cache=cache, visited=visited) for r in inputs]
|
|
|
|
def _update_multi_inputs(self):
|
|
if not self.get('multi_sockets'): return
|
|
for k, ms in self['multi_sockets'].items():
|
|
def is_active(input) -> bool:
|
|
a: bool = input.is_linked or getattr(input, 'value', False)
|
|
b: bool = getattr(inputs()[-1], 'value', False)
|
|
return a or b
|
|
|
|
def inputs():
|
|
return [i for i in list(self.inputs) if i.name == ms['name']]
|
|
|
|
socket_type = ms['bl_idname']
|
|
name = ms['name']
|
|
|
|
if len(inputs()) == 0:
|
|
self.add_input(socket_type, name, name+str(len(inputs())))
|
|
return
|
|
|
|
if is_active(inputs()[-1]):
|
|
self.add_input(socket_type, name, name+str(len(inputs())))
|
|
|
|
while True:
|
|
if len(inputs()) < 2: break
|
|
if not is_active(inputs()[-1]) and not is_active(inputs()[-2]):
|
|
self.inputs.remove(inputs()[-1])
|
|
continue
|
|
break
|
|
|
|
## Evaluation helpers
|
|
|
|
_current_execution_cache: _EvalCache | None = None
|
|
|
|
@contextmanager
|
|
def execution_cache():
|
|
global _current_execution_cache
|
|
if _current_execution_cache is not None:
|
|
yield _current_execution_cache
|
|
return
|
|
|
|
_current_execution_cache = {}
|
|
try:
|
|
yield _current_execution_cache
|
|
finally:
|
|
_current_execution_cache = None
|
|
|
|
@contextmanager
|
|
def node_execution_feedback(node: bpy.types.Node):
|
|
# Save original color
|
|
old_use_color = node.use_custom_color
|
|
old_color = node.color[:]
|
|
|
|
node.use_custom_color = True
|
|
node.color = (75/255, 73/255, 19/255) # Green
|
|
|
|
# Refresh UI
|
|
bpy.ops.wm.redraw_timer(type='DRAW_WIN_SWAP', iterations=1)
|
|
try:
|
|
with execution_cache():
|
|
yield
|
|
finally:
|
|
node.use_custom_color = old_use_color
|
|
node.color = old_color
|
|
bpy.ops.wm.redraw_timer(type='DRAW_WIN_SWAP', iterations=1)
|
|
|
|
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 cache is None:
|
|
cache = getattr(socket.node, "_current_cache", _current_execution_cache)
|
|
if visited is None:
|
|
visited = getattr(socket.node, "_current_visited", None)
|
|
|
|
if not socket.is_linked:
|
|
val = getattr(socket, "value", None)
|
|
if val is not None:
|
|
return val
|
|
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,
|
|
):
|
|
if cache is None:
|
|
cache = _current_execution_cache
|
|
|
|
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, socket.name)
|
|
if key in cache: return cache[key]
|
|
|
|
node_key = (socket.node.as_pointer(), socket.node.name)
|
|
if node_key in visited: return default # cycle guard
|
|
|
|
visited.add(node_key)
|
|
try:
|
|
with node_execution_feedback(socket.node):
|
|
if not isinstance(socket.node, SliceNode):
|
|
return default
|
|
|
|
# Set these temporarily so self.resolve() can find them
|
|
socket.node._current_cache = cache # type: ignore
|
|
socket.node._current_visited = visited # type: ignore
|
|
|
|
try:
|
|
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:
|
|
if hasattr(socket.node, "_current_cache"): del socket.node._current_cache
|
|
if hasattr(socket.node, "_current_visited"): del socket.node._current_visited
|
|
finally:
|
|
visited.remove(node_key)
|