211 lines
7.5 KiB
Python
211 lines
7.5 KiB
Python
from __future__ import annotations
|
|
from typing import Any, Callable
|
|
from bpy.types import NodeSocket
|
|
from .list_utils import is_input_a_list, is_socket_vector_type, prepare_indexed_inputs
|
|
from .node_errors import attach_node_failure_info
|
|
|
|
def evaluate_node_multivalue(
|
|
node: Any,
|
|
raw_evaluate_fn: Callable[..., Any],
|
|
socket_name: str,
|
|
**kwargs: Any
|
|
) -> Any:
|
|
"""Implicitly evaluates raw_evaluate_fn element-wise across list inputs.
|
|
If no inputs are lists, executes raw_evaluate_fn once and returns the result.
|
|
If list lengths mismatch (and are not length 1 / scalar), raises ValueError.
|
|
"""
|
|
# Populate inputs if empty in unit tests
|
|
if hasattr(node, "init") and (not hasattr(node, "inputs") or len(node.inputs) == 0):
|
|
try:
|
|
node.init(None)
|
|
except Exception:
|
|
pass
|
|
|
|
inputs_col = getattr(node, "inputs", None)
|
|
inputs_list = list(inputs_col) if inputs_col is not None else []
|
|
|
|
# Check if instance has a mocked resolve method (e.g. in unit tests)
|
|
is_mocked = hasattr(node, "__dict__") and "resolve" in node.__dict__
|
|
mock_resolver = node.__dict__["resolve"] if is_mocked else None
|
|
|
|
# Step 1: Collect raw values for all input sockets
|
|
input_keys: list[Any] = []
|
|
raw_inputs: dict[Any, Any] = {}
|
|
key_to_name: dict[Any, str] = {}
|
|
name_to_keys: dict[str, list[Any]] = {}
|
|
|
|
for idx, sock in enumerate(inputs_list):
|
|
s_name = getattr(sock, "name", getattr(sock, "identifier", f"input_{idx}"))
|
|
key = sock
|
|
input_keys.append(key)
|
|
key_to_name[key] = s_name
|
|
name_to_keys.setdefault(s_name, []).append(key)
|
|
|
|
# Check list_socket_mapping for any mapped inputs not present in inputs_list
|
|
mapping = getattr(node, "list_socket_mapping", None)
|
|
if mapping:
|
|
for item in mapping:
|
|
m_in = item[0] if isinstance(item, (list, tuple)) and item else None
|
|
if m_in and m_in not in name_to_keys:
|
|
input_keys.append(m_in)
|
|
key_to_name[m_in] = m_in
|
|
name_to_keys.setdefault(m_in, []).append(m_in)
|
|
|
|
def _query_val(key: Any) -> Any:
|
|
if mock_resolver is not None:
|
|
s_name = key_to_name[key]
|
|
# If multiple sockets share the same name (e.g. MathNode), query by socket object first
|
|
if len(name_to_keys.get(s_name, [])) > 1:
|
|
try:
|
|
val = mock_resolver(key)
|
|
if val is not None:
|
|
return val
|
|
except Exception:
|
|
pass
|
|
# Otherwise query by socket name
|
|
try:
|
|
val = mock_resolver(s_name)
|
|
if val is not None:
|
|
return val
|
|
except Exception:
|
|
pass
|
|
try:
|
|
return mock_resolver(key)
|
|
except Exception:
|
|
return None
|
|
else:
|
|
return node.resolve(key)
|
|
|
|
for key in input_keys:
|
|
raw_inputs[key] = _query_val(key)
|
|
|
|
if not input_keys:
|
|
return raw_evaluate_fn(node, socket_name, **kwargs)
|
|
|
|
# Step 2: Determine list status and vector/atomic flags for each input
|
|
always_in = set(getattr(node, "always_list_inputs", []) or [])
|
|
has_any_list = False
|
|
is_vec_map: dict[Any, bool] = {}
|
|
is_list_map: dict[Any, bool] = {}
|
|
|
|
for key in input_keys:
|
|
s_name = key_to_name[key]
|
|
val = raw_inputs[key]
|
|
is_vec = is_socket_vector_type(node, s_name)
|
|
is_vec_map[key] = is_vec
|
|
is_atom = s_name in always_in
|
|
|
|
is_list = is_input_a_list(val, is_vector=is_vec)
|
|
if is_atom:
|
|
# Atomic list inputs (e.g. Layers) are only multivalue if nested
|
|
is_nested = (
|
|
isinstance(val, (list, tuple))
|
|
and len(val) > 0
|
|
and isinstance(val[0], (list, tuple))
|
|
)
|
|
is_list = is_nested
|
|
|
|
is_list_map[key] = is_list
|
|
if is_list:
|
|
has_any_list = True
|
|
|
|
# Fast path: no list inputs
|
|
if not has_any_list:
|
|
return raw_evaluate_fn(node, socket_name, **kwargs)
|
|
|
|
# Step 3: Identify main socket key
|
|
main_key: Any = None
|
|
if mapping and len(mapping) > 0 and mapping[0][0]:
|
|
mapped_name = mapping[0][0]
|
|
matching_keys = name_to_keys.get(mapped_name, [])
|
|
if matching_keys:
|
|
main_key = matching_keys[0]
|
|
|
|
if main_key is None or main_key not in raw_inputs:
|
|
# Pick the first key that is a list
|
|
main_key = next((k for k in input_keys if is_list_map.get(k, False)), input_keys[0])
|
|
|
|
main_val = raw_inputs[main_key]
|
|
main_is_vec = is_vec_map.get(main_key, False)
|
|
main_name = key_to_name[main_key]
|
|
|
|
# Prepare secondary inputs
|
|
sec_keys = [k for k in input_keys if k != main_key]
|
|
sec_names = [key_to_name[k] for k in sec_keys]
|
|
sec_vals = [raw_inputs[k] for k in sec_keys]
|
|
sec_vec_flags = [is_vec_map[k] for k in sec_keys]
|
|
sec_atom_flags = [key_to_name[k] in always_in for k in sec_keys]
|
|
|
|
try:
|
|
main_is_list, main_items, sec_broadcasted = prepare_indexed_inputs(
|
|
main_val,
|
|
*sec_vals,
|
|
is_vector=sec_vec_flags,
|
|
input_names=sec_names,
|
|
is_atomic=sec_atom_flags,
|
|
main_is_vector=main_is_vec,
|
|
)
|
|
except Exception as e:
|
|
attach_node_failure_info(node, e, fallback_inputs=[key_to_name[k] for k in input_keys])
|
|
raise
|
|
|
|
n = len(main_items)
|
|
if n == 0:
|
|
return []
|
|
|
|
# Map broadcasted lists back to keys
|
|
broadcasted_by_key: dict[Any, list[Any]] = {main_key: main_items}
|
|
for idx, sec_k in enumerate(sec_keys):
|
|
broadcasted_by_key[sec_k] = sec_broadcasted[idx]
|
|
|
|
# Step 4: Element-wise execution
|
|
results: list[Any] = []
|
|
node._in_multivalue_eval = True
|
|
|
|
try:
|
|
for i in range(n):
|
|
elem_ctx: dict[Any, Any] = {}
|
|
for key in input_keys:
|
|
b_val = broadcasted_by_key[key][i]
|
|
elem_ctx[key] = b_val
|
|
# Also index by string name and identifier
|
|
s_name = key_to_name[key]
|
|
# If multiple sockets have the same name, the first one set in elem_ctx is default,
|
|
# but socket-object indexing disambiguates.
|
|
if s_name not in elem_ctx:
|
|
elem_ctx[s_name] = b_val
|
|
if hasattr(key, "identifier"):
|
|
elem_ctx[key.identifier] = b_val
|
|
|
|
node._eval_iteration_ctx = elem_ctx
|
|
|
|
if is_mocked:
|
|
# Intercept mocked resolver to return element values
|
|
def _mock_elem_resolve(label, _ctx=elem_ctx, _orig=mock_resolver):
|
|
if label in _ctx:
|
|
return _ctx[label]
|
|
name = getattr(label, "name", None)
|
|
if name in _ctx:
|
|
return _ctx[name]
|
|
ident = getattr(label, "identifier", None)
|
|
if ident in _ctx:
|
|
return _ctx[ident]
|
|
return _orig(label)
|
|
|
|
node.resolve = _mock_elem_resolve
|
|
|
|
item_res = raw_evaluate_fn(node, socket_name, **kwargs)
|
|
results.append(item_res)
|
|
|
|
except Exception as e:
|
|
attach_node_failure_info(node, e, fallback_inputs=[key_to_name[k] for k in input_keys])
|
|
raise
|
|
finally:
|
|
node._in_multivalue_eval = False
|
|
node._eval_iteration_ctx = None
|
|
if is_mocked:
|
|
node.resolve = mock_resolver
|
|
elif hasattr(node, "__dict__") and "resolve" in node.__dict__:
|
|
node.__dict__.pop("resolve", None)
|
|
|
|
return results
|