456 lines
16 KiB
Python
456 lines
16 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, register_class
|
|
class SliceNode(Node, AutoBlIdName):
|
|
def update(self) -> None:
|
|
self._update_multi_groups()
|
|
|
|
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)
|
|
inputs = getattr(self, "inputs", None)
|
|
if inputs is None:
|
|
return None
|
|
return resolve_input(inputs.get(fieldname), cache=cache, visited=visited)
|
|
|
|
def add_output(self, bl_idname: str, name: str = "", identifier: str = "", hide: bool = False, display_shape: str = ""):
|
|
eval_name = name or get_label_by_name(bl_idname)
|
|
sock = self.outputs.new(bl_idname, eval_name, identifier=identifier or eval_name)
|
|
sock.hide = hide
|
|
if display_shape:
|
|
sock.display_shape = display_shape
|
|
return sock
|
|
|
|
def add_input(self, bl_idname: str, name: str = "", identifier: str = "", hide: bool = False, display_shape: str = ""):
|
|
eval_name = name or get_label_by_name(bl_idname)
|
|
sock = self.inputs.new(bl_idname, eval_name, identifier=identifier or eval_name)
|
|
sock.hide = hide
|
|
if display_shape:
|
|
sock.display_shape = display_shape
|
|
return sock
|
|
|
|
def add_multi_socket_group(self, in_type: str, out_type: str = "", name: str = "", in_name: str = "", out_name: str = "", in_display_shape: str = "", out_display_shape: str = ""):
|
|
if not self.get('multi_groups'): self['multi_groups'] = {}
|
|
group_id = name or in_name or out_name or in_type
|
|
self['multi_groups'][group_id] = {
|
|
'in_type': in_type,
|
|
'out_type': out_type,
|
|
'in_name': in_name or name or get_label_by_name(in_type),
|
|
'out_name': out_name or name or (get_label_by_name(out_type) if out_type else ""),
|
|
'in_display_shape': in_display_shape,
|
|
'out_display_shape': out_display_shape,
|
|
}
|
|
self._update_multi_groups()
|
|
|
|
def resolve_multi_input(self, fieldname) -> Any:
|
|
cache = getattr(self, "_current_cache", None)
|
|
visited = getattr(self, "_current_visited", None)
|
|
# Search by name or identifier prefix for multi-groups
|
|
inputs = [i for i in list(self.inputs) if i.name == fieldname or i.identifier.startswith(f"{fieldname}_ms_in_")]
|
|
return [resolve_input(r, cache=cache, visited=visited) for r in inputs]
|
|
|
|
def _update_multi_groups(self):
|
|
if not self.get('multi_groups'): return
|
|
|
|
for group_id, cfg in self['multi_groups'].items():
|
|
in_type = cfg['in_type']
|
|
out_type = cfg.get('out_type', '')
|
|
in_name = cfg['in_name']
|
|
out_name = cfg['out_name']
|
|
in_display_shape = cfg.get('in_display_shape', '')
|
|
out_display_shape = cfg.get('out_display_shape', '')
|
|
|
|
def get_ins(): return [i for i in self.inputs if i.identifier.startswith(f"{group_id}_ms_in_")]
|
|
def get_outs(): return [o for o in self.outputs if o.identifier.startswith(f"{group_id}_ms_out_")]
|
|
def is_active(s): return s.is_linked or bool(getattr(s, 'value', False))
|
|
|
|
ins = get_ins()
|
|
|
|
if not ins:
|
|
self.add_input(in_type, in_name, identifier=f"{group_id}_ms_in_0", display_shape=in_display_shape)
|
|
if out_type:
|
|
self.add_output(out_type, out_name, identifier=f"{group_id}_ms_out_0", display_shape=out_display_shape)
|
|
ins = get_ins()
|
|
|
|
if is_active(ins[-1]):
|
|
idx = len(ins)
|
|
self.add_input(in_type, in_name, identifier=f"{group_id}_ms_in_{idx}", display_shape=in_display_shape)
|
|
if out_type:
|
|
self.add_output(out_type, out_name, identifier=f"{group_id}_ms_out_{idx}", display_shape=out_display_shape)
|
|
ins = get_ins()
|
|
|
|
while len(ins) > 1:
|
|
if not is_active(ins[-1]) and not is_active(ins[-2]):
|
|
self.inputs.remove(ins[-1])
|
|
outs = get_outs()
|
|
if out_type and len(outs) >= len(ins):
|
|
self.outputs.remove(outs[-1])
|
|
ins.pop()
|
|
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)
|
|
|
|
@contextmanager
|
|
def _group_context(node: bpy.types.Node):
|
|
stack = _get_group_stack()
|
|
stack.append(node)
|
|
try:
|
|
yield
|
|
finally:
|
|
stack.pop()
|
|
|
|
def _get_group_stack() -> list[bpy.types.Node]:
|
|
cache = _current_execution_cache
|
|
if cache is None: return []
|
|
if '_group_stack' not in cache:
|
|
cache['_group_stack'] = []
|
|
return cache['_group_stack']
|
|
|
|
def _get_current_group_node() -> bpy.types.Node | None:
|
|
stack = _get_group_stack()
|
|
return stack[-1] if stack else None
|
|
|
|
@register_class
|
|
class SLICENODE_Group(bpy.types.NodeCustomGroup, AutoBlIdName):
|
|
bl_label = "Group"
|
|
bl_icon = 'NODETREE'
|
|
|
|
def _tree_updated(self, context):
|
|
try:
|
|
if not self.name: return
|
|
_ = len(self.inputs)
|
|
_ = len(self.outputs)
|
|
except Exception:
|
|
return
|
|
|
|
self.sync_sockets()
|
|
|
|
slicer_tree: bpy.props.PointerProperty( # pyright: ignore[reportInvalidTypeForm]
|
|
type=bpy.types.NodeTree,
|
|
name="Node Tree",
|
|
update=_tree_updated
|
|
)
|
|
|
|
def sync_sockets(self):
|
|
try:
|
|
if not self.name: return
|
|
_ = len(self.inputs)
|
|
_ = len(self.outputs)
|
|
except Exception:
|
|
return
|
|
|
|
if not self.slicer_tree:
|
|
self.inputs.clear()
|
|
self.outputs.clear()
|
|
return
|
|
|
|
# 1. Collect expected sockets from interface
|
|
expected_inputs = [] # (identifier, name, type)
|
|
expected_outputs = []
|
|
|
|
has_interface = hasattr(self.slicer_tree, "interface")
|
|
if has_interface:
|
|
# In Blender 4.0+, the collection of items is items_tree
|
|
items_tree = getattr(self.slicer_tree.interface, "items_tree", [])
|
|
for item in items_tree:
|
|
if getattr(item, "item_type", "") != 'SOCKET': continue
|
|
target_list = expected_inputs if item.in_out == 'INPUT' else expected_outputs
|
|
target_list.append((item.identifier, item.name, item.socket_type))
|
|
else:
|
|
# Fallback for older Blender API without interface
|
|
for sock in getattr(self.slicer_tree, "inputs", []):
|
|
expected_inputs.append((sock.identifier, sock.name, sock.bl_idname))
|
|
for sock in getattr(self.slicer_tree, "outputs", []):
|
|
expected_outputs.append((sock.identifier, sock.name, sock.bl_idname))
|
|
|
|
# 2. Sync Inputs
|
|
self._sync_socket_collection(self.inputs, expected_inputs)
|
|
|
|
# 3. Sync Outputs
|
|
self._sync_socket_collection(self.outputs, expected_outputs)
|
|
|
|
def _sync_socket_collection(self, collection, expected):
|
|
# Remove sockets that shouldn't be there
|
|
expected_ids = {e[0] for e in expected}
|
|
to_remove = [s for s in collection if s.identifier not in expected_ids]
|
|
|
|
for sock in to_remove:
|
|
collection.remove(sock)
|
|
|
|
# Add or update sockets
|
|
sock_indexes = {s.identifier: i for i, s in enumerate(collection)}
|
|
for ident, name, stype in expected:
|
|
if ident not in sock_indexes:
|
|
try:
|
|
# Explicitly pass identifier to keep it perfectly synced
|
|
sock = collection.new(stype, name, identifier=ident)
|
|
sock_indexes[sock.identifier] = len(collection) - 1
|
|
except Exception as e:
|
|
print(f"Failed to create socket {name} ({stype}): {e}")
|
|
else:
|
|
# Update name if changed
|
|
sock = collection[sock_indexes[ident]]
|
|
if sock.name != name:
|
|
sock.name = name
|
|
|
|
def draw_buttons(self, context, layout):
|
|
assert layout
|
|
row = layout.row(align=True)
|
|
row.prop(self, "slicer_tree", text="")
|
|
if self.slicer_tree:
|
|
op = row.operator("slicenode.group_toggle", text="", icon='NODETREE')
|
|
op.node_name = self.name
|
|
|
|
def draw_label(self):
|
|
if self.slicer_tree:
|
|
return self.slicer_tree.name
|
|
return "Group"
|
|
|
|
|
|
def _resolve_group_output(socket: NodeSocket, cache: _EvalCache | None, visited: set[tuple[int, str]] | None):
|
|
node = socket.node
|
|
assert isinstance(node, (bpy.types.NodeGroup, SLICENODE_Group))
|
|
|
|
if isinstance(node, SLICENODE_Group):
|
|
tree = node.slicer_tree
|
|
else:
|
|
tree = node.node_tree
|
|
|
|
if not tree: return None
|
|
|
|
try:
|
|
idx = list(node.outputs).index(socket)
|
|
except (ValueError, IndexError):
|
|
return None
|
|
|
|
# Find the NodeGroupOutput node that has this input linked, or the first one
|
|
output_node = None
|
|
for n in tree.nodes:
|
|
if isinstance(n, bpy.types.NodeGroupOutput):
|
|
# NodeGroupOutput inputs usually match the group outputs by index
|
|
if idx < len(n.inputs) and n.inputs[idx].is_linked:
|
|
output_node = n
|
|
break
|
|
|
|
if not output_node:
|
|
output_node = next((n for n in tree.nodes if isinstance(n, bpy.types.NodeGroupOutput)), None)
|
|
|
|
if not output_node or idx >= len(output_node.inputs): return None
|
|
|
|
inner_socket = output_node.inputs[idx]
|
|
|
|
with _group_context(node):
|
|
return resolve_input(typing.cast(Any, inner_socket), cache=cache, visited=visited)
|
|
|
|
def _resolve_group_input(socket: NodeSocket, cache: _EvalCache | None, visited: set[tuple[int, str]] | None):
|
|
group_node = _get_current_group_node()
|
|
if not group_node: return None
|
|
|
|
# NodeGroupInput outputs usually match the group inputs by index
|
|
try:
|
|
idx = list(socket.node.outputs).index(socket)
|
|
outer_socket = group_node.inputs[idx]
|
|
except (ValueError, IndexError):
|
|
return None
|
|
|
|
return resolve_input(typing.cast(Any, outer_socket), cache=cache, visited=visited)
|
|
|
|
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
|
|
# Include group stack in cache key to ensure independent evaluation for multiple instances
|
|
group_stack = tuple(n.as_pointer() for n in _get_group_stack())
|
|
|
|
# Use identifier instead of name for the cache key, as multiple sockets
|
|
# (especially in multi-groups) can share the same name but must have unique identifiers.
|
|
key = (group_stack, socket.node.as_pointer(), socket.node.name, socket.identifier)
|
|
if key in cache: return cache[key]
|
|
|
|
node_key = (group_stack, 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):
|
|
if isinstance(socket.node, (bpy.types.NodeGroup, SLICENODE_Group)):
|
|
return _resolve_group_output(socket, cache=cache, visited=visited)
|
|
if isinstance(socket.node, bpy.types.NodeGroupInput):
|
|
return _resolve_group_input(socket, cache=cache, visited=visited)
|
|
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,
|
|
identifier=socket.identifier,
|
|
socket=socket,
|
|
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)
|