procfunc 0.30.0__py3-none-any.whl

procfunc/nodes/execute/infer_runtime_data_type.py

@@ -0,0 +1,195 @@
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import bpy
+import numpy as np
+
+from procfunc import compute_graph as cg
+from procfunc.nodes import bpy_node_info as bni
+from procfunc.nodes import types as nt
+from procfunc.nodes.bindings_util import (
+    RuntimeResolveDataType,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class VectorLike:
+    # named/typed version of None to be used with _infer_value_math_type
+    pass
+
+
+def _infer_value_math_type(
+    val: bpy.types.NodeSocket | Any,
+    py_input: cg.Node | Any,
+    coerce_integers: bool = False,
+) -> bni.NodeDataType | VectorLike | None:
+    """
+    What data type of math should we do on this `val`?
+
+    Args:
+        val: The value to infer the data type of
+        coerce_integers: Whether to coerce integers to floats. Usually used only for shader context, since bl4.2 shaders dont support integer math
+
+    Returns:
+        bni.NodeDataType: The data type of math to do on this `val`
+        VectorLike: The argument is vaguely vector-like, but we leave it up to other vals to decide between Vector Color etc.
+    """
+
+    assert not isinstance(py_input, nt.ProcNode), py_input
+
+    logger.debug(
+        f"{_infer_value_math_type.__name__} starting for {type(val)=} {py_input=} {type(py_input)=}"
+    )
+
+    if (
+        isinstance(py_input, cg.Node)
+        and (vt := py_input.metadata.get("known_value_type", None)) is not None
+    ):
+        socket_type = bni.PYTHON_TYPE_TO_SOCKET_TYPE[vt]
+        res = bni.SOCKET_CLASS_TO_DATATYPE[socket_type.value]
+        logger.debug(f"used known_value_type={vt} from {py_input=} to infer {res=}")
+    elif isinstance(val, bpy.types.NodeSocket):
+        res = bni.SOCKET_DTYPE_TO_DATATYPE[bni.SocketDType(val.type)]
+        logger.debug(f"used {val.type=} to infer {res=}")
+    elif type(val) in bni.PYTHON_TYPE_TO_SOCKET_TYPE:
+        res = bni.PYTHON_TYPE_TO_SOCKET_TYPE[type(val)]
+        res = bni.SOCKET_CLASS_TO_DATATYPE[res.value]
+    elif isinstance(val, (list, tuple, np.ndarray)):
+        return VectorLike()
+    elif val is None:
+        return None
+    else:
+        raise ValueError(
+            f"{_infer_value_math_type.__name__} got {val=} of type {type(val)} "
+            f"which is not handled by either {bni.PYTHON_TYPE_TO_SOCKET_TYPE.keys()} "
+            f"or {bni.SOCKET_DTYPE_TO_DATATYPE.keys()}"
+        )
+
+    if coerce_integers and res == bni.NodeDataType.INT:
+        res = bni.NodeDataType.FLOAT
+
+    return res
+
+
+_vectorlike_types = {bni.NodeDataType.FLOAT_VECTOR, bni.NodeDataType.RGBA}
+
+
+def infer_operation_type(
+    node: cg.Node,
+    inputs: dict[str | int, Any],
+    coerce_integers: bool,
+    filter_keys: list[str] | None = None,
+    vectorlike_default: bni.NodeDataType | None = None,
+) -> bni.NodeDataType:
+    filtered_keys = [
+        k for k in node.kwargs.keys() if filter_keys is None or k in filter_keys
+    ]
+    input_data_types: list[bni.NodeDataType | VectorLike] = [
+        _infer_value_math_type(inputs[k], node.kwargs[k], coerce_integers)
+        for k in filtered_keys
+    ]
+
+    input_data_types.extend(
+        [
+            _infer_value_math_type(inputs[i], arg, coerce_integers)
+            for i, arg in enumerate(node.args)
+        ]
+    )
+
+    specific_types = [
+        v for v in input_data_types if not isinstance(v, VectorLike) and v is not None
+    ]
+
+    if len(specific_types) == 0 and len(input_data_types) > 0:
+        if vectorlike_default is not None:
+            return vectorlike_default
+        # we had ALL VectorLike
+        raise NotImplementedError(
+            f"Need to handle case where all arguments to {node=} are non-type-specific tuples/lists. Potentially just assume Vector?"
+        )
+
+    data_type = specific_types[0]
+
+    if any(v != data_type for v in specific_types[1:]):
+        dtypes_msg = " ".join([str(v) for v in input_data_types])
+        raise ValueError(
+            f"Attempted to create operator with non-matching input types {dtypes_msg}. "
+            f"Need to use .astype to hint to make arg types match a single operation datatype, e.g. .astype(float) or .astype(pf.Color)"
+        )
+
+    if (
+        len(specific_types) > 0
+        and any(isinstance(v, VectorLike) for v in specific_types)
+        and data_type not in [bni.NodeDataType.FLOAT_VECTOR, bni.NodeDataType.RGBA]
+    ):
+        raise ValueError(
+            f"{input_data_types=} has compatibile types EXCEPT that {data_type=} "
+            "is not compatible with unspecified tuples/lists (shown as VectorLike)."
+        )
+
+    logger.debug(f"inferred {data_type=} for {node=}")
+
+    return data_type
+
+
+_vectorlike_types = {bni.NodeDataType.FLOAT_VECTOR, bni.NodeDataType.RGBA}
+
+
+def resolve_operation_data_type(
+    node: cg.Node,
+    input_results: dict[str | int, Any],
+    resolve_options: RuntimeResolveDataType,
+    coerce_integers: bool,
+) -> bni.NodeDataType:
+    vectorlike_options = [
+        o for o in resolve_options.data_types if o in _vectorlike_types
+    ]
+    vectorlike_default = vectorlike_options[0] if len(vectorlike_options) > 0 else None
+
+    data_type = infer_operation_type(
+        node,
+        input_results,
+        coerce_integers,
+        filter_keys=resolve_options.dependent_input_names,
+        vectorlike_default=vectorlike_default,
+    )
+
+    if logger.isEnabledFor(logging.DEBUG):
+        logger.debug(
+            f"{resolve_operation_data_type.__name__} for {node=} inferred {data_type=}"
+        )
+
+    return data_type
+
+
+def map_data_type_for_differing_node_interface(
+    data_type: bni.NodeDataType,
+    bl_node: bpy.types.Node,
+    attr_key: str,
+) -> str:
+    """
+    As of bl4.2, some nodes use differing sets of strings to specify data types.Any
+
+    We look at the options to pick which convention to UserWarning
+
+    NOTE: this may one day become unnecessary in future blender versions,
+        but this function will end up silently always doing the same option
+    """
+
+    data_type_options: list[str] = list(
+        bl_node.bl_rna.properties[attr_key].enum_items.keys()
+    )
+    if data_type.value in data_type_options:
+        # MapRange seems to use the DataType naming convention
+        return data_type.value
+    as_socket_dtype = bni.DATATYPE_TO_SOCKET_DTYPE[data_type]
+    if as_socket_dtype.value in data_type_options:
+        # wheras others e.g. Mix use the SocketType naming convention
+        return as_socket_dtype.value
+    raise ValueError(
+        f"Failed resolve {data_type=} or {as_socket_dtype=} to an available {data_type_options=} "
+        f"for {bl_node=} {attr_key=}"
+    )

procfunc/nodes/execute/util.py

@@ -0,0 +1,247 @@
+from dataclasses import dataclass
+from typing import Any, Callable
+
+import bpy
+
+from procfunc import compute_graph as cg
+from procfunc import types as t
+from procfunc.nodes import bpy_node_info as bni
+from procfunc.nodes import func, math
+from procfunc.nodes import types as nt
+from procfunc.util.log import add_exception_context_msg
+
+
+@dataclass
+class NodeOperatorResolution:
+    pf_func: Callable[..., Any]
+    value_type: bni.NodeDataType
+    operator_type: cg.OperatorType
+
+
+def _float_math_defs() -> list[NodeOperatorResolution]:
+    return [
+        NodeOperatorResolution(math.add, bni.NodeDataType.FLOAT, cg.OperatorType.ADD),
+        NodeOperatorResolution(
+            math.subtract, bni.NodeDataType.FLOAT, cg.OperatorType.SUB
+        ),
+        NodeOperatorResolution(
+            math.multiply, bni.NodeDataType.FLOAT, cg.OperatorType.MUL
+        ),
+        NodeOperatorResolution(
+            math.divide, bni.NodeDataType.FLOAT, cg.OperatorType.DIV
+        ),
+        NodeOperatorResolution(math.power, bni.NodeDataType.FLOAT, cg.OperatorType.POW),
+        NodeOperatorResolution(
+            math.modulo, bni.NodeDataType.FLOAT, cg.OperatorType.MOD
+        ),
+        NodeOperatorResolution(
+            math.greater_than, bni.NodeDataType.FLOAT, cg.OperatorType.GREATER_THAN
+        ),
+        NodeOperatorResolution(
+            math.less_than, bni.NodeDataType.FLOAT, cg.OperatorType.LESS_THAN
+        ),
+    ]
+
+
+def _vector_math_defs(node_data_type: bni.NodeDataType) -> list[NodeOperatorResolution]:
+    return [
+        NodeOperatorResolution(math.vector_add, node_data_type, cg.OperatorType.ADD),
+        NodeOperatorResolution(
+            math.vector_subtract, node_data_type, cg.OperatorType.SUB
+        ),
+        NodeOperatorResolution(
+            math.vector_multiply, node_data_type, cg.OperatorType.MUL
+        ),
+        NodeOperatorResolution(math.vector_divide, node_data_type, cg.OperatorType.DIV),
+    ]
+
+
+NODE_OPERATOR_TABLE = [
+    *_float_math_defs(),
+    *_vector_math_defs(bni.NodeDataType.FLOAT_VECTOR),
+    # *_vector_math_defs(NodeDataType.RGBA), # no longer used, instead we will require explicit conversion to vector before math
+    NodeOperatorResolution(
+        math.vector_modulo, bni.NodeDataType.FLOAT_VECTOR, cg.OperatorType.MOD
+    ),
+    # NodeOperatorResolution(
+    #     func.combine_xyz, bni.NodeDataType.FLOAT_VECTOR, cg.OperatorType.VECTOR_PACK
+    # ),
+    NodeOperatorResolution(
+        func.separate_xyz, bni.NodeDataType.FLOAT_VECTOR, cg.OperatorType.NOOP
+    ),
+]
+
+
+def normalize_socket_type(socket_type: str) -> str:
+    # Map all vector subtypes to basic NodeSocketVector
+    if socket_type.startswith("NodeSocketVector"):
+        return "NodeSocketVector"
+    elif socket_type.startswith("NodeSocketFloat"):
+        return "NodeSocketFloat"
+    elif socket_type.startswith("NodeSocketInt"):
+        return "NodeSocketInt"
+    elif socket_type == "NodeSocketVirtual":
+        raise NotImplementedError(
+            f"Virtual sockets are not supported, got {socket_type}"
+        )
+    else:
+        return socket_type
+
+
+def infer_socket_type_from_value(input_val: Any) -> str:
+    """Infer the appropriate Blender socket type from an input value."""
+    match input_val:
+        case bpy.types.NodeSocket():
+            socket_type = (
+                input_val.bl_idname
+                if hasattr(input_val, "bl_idname")
+                else input_val.__class__.__name__
+            )
+            return normalize_socket_type(socket_type)
+        case x if isinstance(x, bpy.types.NodeInternal):
+            try:
+                first_enabled = next(o for o in x.outputs if o.enabled)
+                return infer_socket_type_from_value(first_enabled)
+            except StopIteration:
+                return "NodeSocketFloat"
+        case int():
+            return "NodeSocketInt"
+        case float():
+            return "NodeSocketFloat"
+        case bool():
+            return "NodeSocketBool"
+        case str():
+            return "NodeSocketString"
+        case x if hasattr(x, "__len__"):
+            match len(x):
+                case 3:
+                    return "NodeSocketVector"
+                case 4:
+                    return "NodeSocketColor"
+                case 1 | 2:
+                    return "NodeSocketFloat"
+                case _:
+                    return "NodeSocketFloat"
+        case _:
+            return "NodeSocketFloat"
+
+
+def get_active_sockets(
+    sockets: bpy.types.bpy_prop_collection,
+) -> list[bpy.types.NodeSocket]:
+    if len(set(s.is_output for s in sockets)) > 1:
+        raise ValueError(f"{sockets=} had a mix of input and output sockets")
+
+    return [
+        socket
+        for socket in sockets
+        if (
+            socket.enabled
+            and socket.identifier != "__extend__"
+            and socket.name != ""  # empty-name sockets are internal to nodegroup wiring
+        )
+    ]
+
+
+def get_nth_socket(
+    sockets: bpy.types.bpy_prop_collection,  # of bpy.types.NodeSocket
+    socket_name: str,
+    index: int,
+    debug_node_name: str = "",
+) -> bpy.types.NodeSocket:
+    """
+    Get the nth occurrence of a socket with the given name
+
+    """
+
+    socket_name_fuzzy = socket_name.lower()
+    socket_spaces = socket_name_fuzzy.replace("_", " ")
+
+    count = 0
+    for socket in sockets:
+        name_lower = socket.name.lower()
+        if socket.enabled and (
+            name_lower == socket_spaces or name_lower == socket_name_fuzzy
+        ):
+            if count == index:
+                return socket
+            count += 1
+
+    enabled = [s.name for s in sockets if s.enabled]
+    disabled = [s.name for s in sockets if not s.enabled]
+
+    if any(d.lower() == socket_spaces for d in disabled):
+        raise ValueError(
+            f"User attempted to use input {socket_name.lower()!r} for node {debug_node_name} but it is disabled. "
+            "We may have failed to set the node's `data_type` or `mode` input arguments"
+        )
+
+    raise ValueError(
+        f"Node input {debug_node_name} has no enabled output socket "
+        f"which fuzzmatches {socket_name=} {index=}, sockets were {enabled=} and {disabled=}"
+    )
+
+
+def get_input_socket_to_connect_to(
+    node_tree: bpy.types.NodeTree,
+    node: bpy.types.Node,
+    socket_id: tuple[str, int] | str,
+    input_val: bpy.types.NodeSocket | bpy.types.Node | Any | None,
+) -> bpy.types.NodeSocket:
+    assert isinstance(node_tree, bpy.types.NodeTree), node_tree
+
+    if isinstance(socket_id, tuple):
+        socket_name, socket_index = socket_id
+    elif isinstance(socket_id, str):
+        socket_name = socket_id
+        socket_index = 0
+    else:
+        raise ValueError(f"Invalid socket id {socket_id=}")
+
+    socket = get_nth_socket(node.inputs, socket_name, socket_index, node.name)
+    assert socket is not None, f"Node {node.name=} has no input {socket_name=}"
+
+    return socket
+
+
+def assign_default_value(
+    target_socket: bpy.types.NodeSocket,
+    input_val: Any,
+    data_type: bni.NodeDataType | None = None,
+):
+    if input_val is None:
+        return
+
+    if data_type is None:
+        data_type = bni.SOCKET_DTYPE_TO_DATATYPE[bni.SocketDType(target_socket.type)]
+
+    assert not isinstance(input_val, nt.ProcNode), input_val
+    assert isinstance(data_type, bni.NodeDataType), data_type
+
+    match data_type:
+        case bni.NodeDataType.RGBA:
+            if len(input_val) == 3:
+                input_val = (input_val[0], input_val[1], input_val[2], 1.0)
+            assert not any(isinstance(x, nt.ProcNode) for x in input_val), input_val
+            target_socket.default_value = input_val
+        case bni.NodeDataType.FLOAT_VECTOR:
+            if hasattr(input_val, "__len__") and len(input_val) == 3:
+                input_val = (input_val[0], input_val[1], input_val[2])
+            assert not any(isinstance(x, nt.ProcNode) for x in input_val), input_val
+            target_socket.default_value = t.Vector(input_val)
+        case bni.NodeDataType.FLOAT:
+            target_socket.default_value = float(input_val)
+        case (
+            bni.NodeDataType.OBJECT
+            | bni.NodeDataType.MATERIAL
+            | bni.NodeDataType.COLLECTION
+        ):
+            if isinstance(input_val, (t.MeshObject, t.Material, t.Collection)):
+                target_socket.default_value = input_val.item()
+            else:
+                target_socket.default_value = input_val
+        case _:
+            with add_exception_context_msg(
+                prefix=f"While assigning {input_val=} to {target_socket.name=} with specified {data_type=}"
+            ):
+                target_socket.default_value = input_val