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