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 _feedback_enabled: bool = True @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 feedback_disabled(): global _feedback_enabled old = _feedback_enabled _feedback_enabled = False try: yield finally: _feedback_enabled = old @contextmanager def node_execution_feedback(node: bpy.types.Node): if not _feedback_enabled: with execution_cache(): yield return # 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", None) or _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: socket.node._current_cache = None # type: ignore socket.node._current_visited = None # type: ignore finally: visited.remove(node_key)