procfunc 0.30.0__py3-none-any.whl

procfunc/__init__.py

@@ -0,0 +1,87 @@
+# ruff: noqa: I001, F401
+# ensure this gets imported first so that mathutils etc is available even if later modules dont import bpy
+import bpy
+
+__version__ = "0.30.0"
+
+from numpy.random import Generator as RNG
+
+# ensure these are always imported first and in the right order
+import bpy as _bpy
+import mathutils as _mu
+
+from . import (
+    compute_graph,
+    control,
+    nodes,
+    ops,
+    random,
+    tracer,
+    util,
+    context,
+    color,
+    transforms,
+)
+from .types import (
+    Vector,
+    Color,
+    Euler,
+    Quaternion,
+    Matrix,
+    BVHTree,
+    Object,
+    CameraObject,
+    MeshObject,
+    CurveObject,
+    VolumeObject,
+    EmptyObject,
+    ArmatureObject,
+    HairObject,
+    LatticeObject,
+    LightObject,
+    LightProbeObject,
+    MetaObject,
+    Material,
+    Texture,
+    Collection,
+    Image,
+    ViewLayer,
+    Asset,
+    World,
+)
+from .nodes import ProcNode, Shader, NodeDataType
+from .tracer import trace, autowrap_module, add_search_scope
+from .util.manifest import module_path
+
+autowrap_module(random)
+autowrap_module(color)
+add_search_scope(nodes)
+add_search_scope(ops)
+
+__all__ = [
+    # Subpackages
+    "compute_graph",
+    "control",
+    "context",
+    "color",
+    "nodes",
+    "ops",
+    "random",
+    "tracer",
+    "transforms",
+    "util",
+    # Node graph primitives
+    "ProcNode",
+    "Shader",
+    "NodeDataType",
+    # Tracing entrypoints
+    "trace",
+    # Utilities
+    "RNG",
+    "module_path",
+    # NOTE: Blender wrapper types (Object, MeshObject, Material, Texture, ...)
+    # are re-exported above for convenience but live in procfunc.types — see
+    # that page for documentation. Blender/mathutils re-exports (Vector, Color,
+    # Euler, Quaternion, Matrix, BVHTree, NodeGroup, Scene, ViewLayer) are
+    # likewise kept importable but documented upstream by Blender.
+]

procfunc/color.py

@@ -0,0 +1,57 @@
+import numpy as np
+
+from procfunc import types as t
+
+
+def hsv_color(
+    hsv: np.ndarray | tuple | None = None,
+    *,
+    hue: float | None = None,
+    saturation: float | None = None,
+    value: float | None = None,
+) -> t.Color:
+    color = t.Color()
+    if hsv is not None:
+        color.hsv = hsv
+    else:
+        color.hsv = (hue, saturation, value)
+    return color
+
+
+hsv_to_rgba = hsv_color
+
+
+def rgb_color(
+    r: float | None = None,
+    g: float | None = None,
+    b: float | None = None,
+    rgb: np.ndarray | None = None,
+) -> t.Color:
+    color = t.Color()
+    if rgb is not None:
+        color.r, color.g, color.b = rgb
+    else:
+        color.r, color.g, color.b = r, g, b
+    return color
+
+
+def _srgb_to_linearrgb(c):
+    if c < 0:
+        return 0
+    elif c < 0.04045:
+        return c / 12.92
+    else:
+        return ((c + 0.055) / 1.055) ** 2.4
+
+
+def _hex_to_rgb(h: int, alpha: float = 1):
+    r = (h & 0xFF0000) >> 16
+    g = (h & 0x00FF00) >> 8
+    b = h & 0x0000FF
+    return tuple([_srgb_to_linearrgb(c / 0xFF) for c in (r, g, b)])
+
+
+def hex_color(h: int, alpha: float = 1):
+    c = t.Color()
+    c.r, c.g, c.b = _hex_to_rgb(h, alpha)
+    return c

procfunc/compute_graph/__init__.py

@@ -0,0 +1,28 @@
+from .compute_graph import (
+    ComputeGraph,
+)
+from .node import (
+    ConstantNode,
+    FunctionCallNode,
+    GetAttributeNode,
+    InputPlaceholderNode,
+    MethodCallNode,
+    MutatedArgumentNode,
+    Node,
+    ProceduralNode,
+    SubgraphCallNode,
+    normalize_args_to_kwargs,
+)
+from .operators_info import OperatorType
+from .proxy import AttributeProxy, Proxy
+from .util import (
+    LiteralConstant,
+    graph_nodes_equal,
+    transform_compute_graph,
+    transform_nodetree,
+    traverse_breadth_first,
+    traverse_depth_first,
+    traverse_depth_first_node,
+    traverse_nested_graphs,
+    usages_per_node,
+)

procfunc/compute_graph/compute_graph.py

@@ -0,0 +1,115 @@
+import itertools
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+from procfunc.util.pytree import PyTree
+
+from .node import Node
+
+logger = logging.getLogger(__name__)
+
+
+def _evaluate_node(node: Node) -> Any:
+    if node.result is not None:
+        return node.result
+
+    for arg in itertools.chain(node.args, node.kwargs.values()):
+        if isinstance(arg, Node):
+            arg.result = _evaluate_node(arg)
+
+    arg_eval, kwarg_eval = node.inputs.map(
+        lambda x: x.result if isinstance(x, Node) else x
+    ).obj()
+
+    # match node:
+    #     case cg.FunctionCallNode:
+    #         return node.func(*arg_vals, **kwarg_vals)
+    #     case cg.PlaceholderNode:
+    #         raise NotImplementedError(
+    #             f"Placeholder {node!r} or '<unnamed>'} should not be evaluated - "
+    #             "its .result should be populated in advance"
+    #         )
+    #     case cg.MutatedArgumentNode:
+    #         # Evaluate the mutation call first (for side effects), then return the original object
+    #         mutation_call_node = node.args[1]
+    #         _evaluate_node(mutation_call_node)
+    #         return (
+    #             node.args[0].result if isinstance(node.args[0], Node) else node.args[0]
+    #         )
+    #     case _:
+    #         raise NotImplementedError(f"Unsupported node operation: {node.kind}")
+
+
+def _clear_node_results(node: Node):
+    node.result = None
+    for arg in itertools.chain(node.args, node.kwargs.values()):
+        if isinstance(arg, Node) and arg.result is not None:
+            _clear_node_results(arg)
+
+
+@dataclass
+class ComputeGraph:
+    inputs: PyTree[Any, Node]
+    outputs: PyTree[Any, Node]
+    name: str
+    metadata: dict[str, Any]
+
+    def __post_init__(self):
+        # input_names = set(self.inputs.names())
+        # if len(input_names) != len(set(input_names)):
+        #     raise ValueError(f"Input names had duplicates: {input_names}")
+
+        # output_names = set(self.outputs.names())
+        # if len(output_names) != len(set(output_names)):
+        #     raise ValueError(f"Output names had duplicates: {output_names}")
+        pass
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.name!r})"
+
+    def clear_values(self):
+        for node in self.outputs.values():
+            _clear_node_results(node)
+
+    def __call__(
+        self,
+        *args,
+        allow_clear: bool = True,
+        **kwargs,
+    ):
+        """
+        Execute the compute graph. If this graph came from a tracer,
+        this should be exactly equivelant to executing the original python function.
+        """
+
+        raise NotImplementedError("Not implemented")
+
+        if len(args) != len(self.inputs):
+            raise ValueError(f"Expected {len(self.inputs)} arguments, got {len(args)}")
+
+        if allow_clear:
+            self.clear_values()
+
+        extra_kwargs = self.kwarg_nodes.keys() - kwargs.keys()
+        if extra_kwargs:
+            raise ValueError(
+                f"{self.__class__.__name__} {self.name!r} got unexpected keyword arguments: {extra_kwargs}"
+            )
+
+        missing_kwargs = kwargs.keys() - self.kwarg_nodes.keys()
+        if missing_kwargs:
+            raise ValueError(
+                f"{self.__class__.__name__} {self.name!r} had missing keyword arguments: {missing_kwargs}"
+            )
+
+        for arg, arg_node in zip(args, self.arg_nodes):
+            arg_node.result = arg
+
+        for k, v in kwargs.items():
+            self.kwarg_nodes[k].result = v
+
+        for node in self.outputs.values():
+            _evaluate_node(node)
+
+        return {name: node.result for name, node in self.outputs.items()}

procfunc/compute_graph/node.py

@@ -0,0 +1,200 @@
+import inspect
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
+
+if TYPE_CHECKING:
+    from procfunc.compute_graph.compute_graph import ComputeGraph
+
+from procfunc.util.pytree import PyTree
+
+logger = logging.getLogger(__name__)
+
+
+class Node:
+    def __init__(self, args: tuple, kwargs: dict, metadata: dict[str, Any] = None):
+        assert isinstance(args, tuple), args
+        assert isinstance(kwargs, dict), kwargs
+        self.args = args
+        self.kwargs = kwargs
+        if metadata is None:
+            metadata = {}
+        self.metadata = metadata
+
+    def inputs_pytree(self) -> PyTree:
+        return PyTree((self.args, self.kwargs))
+
+
+class SubgraphCallNode(Node):
+    def __init__(
+        self,
+        subgraph: "ComputeGraph",
+        args: tuple,
+        kwargs: dict,
+        metadata: dict[str, Any] = None,
+    ):
+        super().__init__(args, kwargs, metadata)
+        self.subgraph = subgraph
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.subgraph.name}, ...)"
+
+
+class FunctionCallNode(Node):
+    def __init__(
+        self,
+        func: Callable[..., Any],
+        args: tuple,
+        kwargs: dict,
+        metadata: dict[str, Any] = None,
+    ):
+        super().__init__(args=args, kwargs=kwargs, metadata=metadata)
+        self.func = func
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.func.__name__}, ...)"
+
+
+class MethodCallNode(Node):
+    """
+    represents an {args[0]}.{method_name}(*args[1:], **kwargs) call
+
+    - the node to be used as `self` is the first arg, since it is a dynamic value
+    - the method name is assumed to be const
+    """
+
+    def __init__(
+        self,
+        callee: Node,
+        method_name: str,
+        args: tuple,
+        kwargs: dict,
+        metadata: dict[str, Any] = None,
+    ):
+        super().__init__(args=(callee, *args), kwargs=kwargs, metadata=metadata)
+        self.method_name = method_name
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.method_name}, ...)"
+
+
+@dataclass
+class GetAttributeNode(Node):
+    def __init__(
+        self, source: Node, attribute_name: str, metadata: dict[str, Any] = None
+    ):
+        # store source as args since it is a Node and may need to be recursively constructed
+        super().__init__(args=(source,), kwargs={}, metadata=metadata)
+        self.attribute_name = attribute_name
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.attribute_name})"
+
+
+@dataclass
+class ProceduralNode(Node):
+    def __init__(
+        self,
+        node_type: str,
+        attrs: dict[str, Any],
+        kwargs: dict,
+        metadata: dict[str, Any] = None,
+    ):
+        super().__init__(args=(), kwargs=kwargs, metadata=metadata)
+        self.node_type = node_type
+
+        for k, v in attrs.items():
+            if isinstance(v, Node):
+                raise ValueError(
+                    f"{self.__class__.__name__}({node_type=}) recieved attrs with non-constant value {k}={v}. "
+                    "(Node values are not allowed as attrs)"
+                )
+        self.attrs = attrs
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.node_type}, ...)"
+
+
+class MutatedArgumentNode(Node):
+    def __init__(
+        self,
+        original_node: "Node",
+        mutator_call_node: "FunctionCallNode | MethodCallNode",
+        metadata: dict[str, Any] = None,
+    ):
+        # store orig/mutator as args() since they are Node and may need to be recursively constructed
+        super().__init__(
+            args=(original_node, mutator_call_node), kwargs={}, metadata=metadata
+        )
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(...)"
+
+
+class ConstantNode(Node):
+    def __init__(self, value: Any, metadata: dict[str, Any] = None):
+        super().__init__(args=(), kwargs={}, metadata=metadata)
+        self.value = value
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.value})"
+
+
+class InputPlaceholderNode(Node):
+    def __init__(self, name: str, default_value: Any, metadata: dict[str, Any] = None):
+        super().__init__(args=(), kwargs={}, metadata=metadata)
+        self.input_name = name
+        self.default_value = default_value
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.default_value})"
+
+
+T = TypeVar("T")
+
+
+def normalize_args_to_kwargs(
+    func: Callable,
+    args: tuple,
+    kwargs: dict,
+) -> tuple[tuple, dict]:
+    """
+    Try to fully populate kwargs, by moving over positional args & filling in defaults
+
+    Some args may not be able to be converted to kwargs, e.g. *args have no names that work
+
+    Args:
+        func: The function whose signature we should respect
+        args: The original positional arguments to the function
+        kwargs: The keyword arguments to the function
+
+    Returns:
+        A tuple of (args, kwargs) where args is a tuple of positional arguments and kwargs is a dictionary of keyword arguments.
+
+    GUARANTEE: func(*returned_args, **returned_kwargs) == func(*args, **kwargs) and does not crash
+    """
+
+    sig = inspect.signature(func)
+    bound = sig.bind(*args, **kwargs)
+    bound.apply_defaults()  # optional: fills in default values
+
+    # Extract any *args parameter back to args tuple
+    remaining_args = ()
+    updated_kwargs = {}
+
+    for param_name, value in bound.arguments.items():
+        param = sig.parameters[param_name]
+        if param.kind == inspect.Parameter.VAR_POSITIONAL:  # *args
+            remaining_args = value
+        elif param.kind == inspect.Parameter.VAR_KEYWORD:  # **kwargs
+            # Unpack **kwargs back to individual kwargs instead of wrapping in dict
+            updated_kwargs.update(value)
+        else:
+            updated_kwargs[param_name] = value
+
+    if False and logger.isEnabledFor(logging.DEBUG):
+        logger.debug(
+            f"normalized {func.__name__} with {args=} {kwargs=} to {remaining_args=} {updated_kwargs=}"
+        )
+
+    return remaining_args, updated_kwargs

procfunc/compute_graph/operators_info.py

@@ -0,0 +1,92 @@
+import operator
+from enum import Enum
+
+
+class OperatorType(Enum):
+    ADD = "add"
+    SUB = "sub"
+    MUL = "mul"
+    DIV = "div"
+    TRUEDIV = "truediv"
+    POW = "pow"
+    MOD = "mod"
+    LESS_THAN = "lt"
+    LESS_THAN_EQUAL = "le"
+    GREATER_THAN = "gt"
+    GREATER_THAN_EQUAL = "ge"
+    EQUAL = "eq"
+    NOT_EQUAL = "ne"
+    AND = "and"
+    OR = "or"
+    NOT = "invert"
+    LSHIFT = "lshift"
+    RSHIFT = "rshift"
+    BIT_AND = "and"
+    BIT_OR = "or"
+    BIT_XOR = "xor"
+
+    GETITEM = "getitem"
+
+    VECTOR_PACK = "VECTOR_PACK"
+    NOOP = "NOOP"
+
+
+OPERATOR_TEMPLATES = {
+    OperatorType.ADD: "{} + {}",
+    OperatorType.SUB: "{} - {}",
+    OperatorType.MUL: "{} * {}",
+    OperatorType.DIV: "{} / {}",
+    OperatorType.POW: "{} ** {}",
+    OperatorType.MOD: "{} % {}",
+    OperatorType.LESS_THAN: "{} < {}",
+    OperatorType.LESS_THAN_EQUAL: "{} <= {}",
+    OperatorType.GREATER_THAN: "{} > {}",
+    OperatorType.GREATER_THAN_EQUAL: "{} >= {}",
+    OperatorType.EQUAL: "{} == {}",
+    OperatorType.NOT_EQUAL: "{} != {}",
+    OperatorType.VECTOR_PACK: "({}, {}, {})",
+    OperatorType.NOOP: "{}",
+    OperatorType.LSHIFT: "{} << {}",
+    OperatorType.RSHIFT: "{} >> {}",
+    OperatorType.BIT_AND: "{} & {}",
+    OperatorType.BIT_OR: "{} | {}",
+    OperatorType.BIT_XOR: "{} ^ {}",
+    OperatorType.GETITEM: "{}[{}]",
+}
+
+OPERATORS_TO_FUNCTIONS = {
+    OperatorType.ADD: operator.add,
+    OperatorType.SUB: operator.sub,
+    OperatorType.MUL: operator.mul,
+    OperatorType.DIV: operator.truediv,
+    OperatorType.MOD: operator.mod,
+    OperatorType.POW: operator.pow,
+    OperatorType.LESS_THAN: operator.lt,
+    OperatorType.LESS_THAN_EQUAL: operator.le,
+    OperatorType.GREATER_THAN: operator.gt,
+    OperatorType.GREATER_THAN_EQUAL: operator.ge,
+    OperatorType.EQUAL: operator.eq,
+    OperatorType.NOT_EQUAL: operator.ne,
+    OperatorType.LSHIFT: operator.lshift,
+    OperatorType.RSHIFT: operator.rshift,
+    OperatorType.BIT_AND: operator.and_,
+    OperatorType.BIT_OR: operator.or_,
+    OperatorType.BIT_XOR: operator.xor,
+    OperatorType.GETITEM: operator.getitem,
+}
+
+FUNCTIONS_TO_OPERATORS = {v: k for k, v in OPERATORS_TO_FUNCTIONS.items()}
+
+REFLECTABLE_OPERATORS = {
+    OperatorType.ADD,
+    OperatorType.SUB,
+    OperatorType.MUL,
+    OperatorType.DIV,
+    OperatorType.MOD,
+    OperatorType.POW,
+    OperatorType.LSHIFT,
+    OperatorType.RSHIFT,
+    OperatorType.BIT_AND,
+    OperatorType.BIT_OR,
+    OperatorType.BIT_XOR,
+}