procfunc 0.30.0__py3-none-any.whl

procfunc/nodes/types.py

@@ -0,0 +1,347 @@
+import copy
+import inspect
+import logging
+from pathlib import Path
+from typing import Any, Generic, TypeVar, Union
+
+import bpy
+
+from procfunc import compute_graph as cg
+from procfunc import types as pt
+from procfunc.compute_graph.operators_info import (
+    OPERATORS_TO_FUNCTIONS,
+    OperatorType,
+)
+from procfunc.util import pytree
+from procfunc.util.manifest import module_path
+
+logger = logging.getLogger(__name__)
+
+INPUT_NODE_TYPE = "NodeGroupInput"
+OUTPUT_NODE_TYPE = "NodeGroupOutput"
+NODE_FUNCTION_INSTANCE_TYPE = "NodeFunctionInstance"
+
+T = TypeVar("T")
+
+
+PROCNODE_OPERATORS = {
+    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,
+}
+
+
+def _node_definition_metadata() -> tuple[str, int, str]:
+    """
+    Dig through this functions callstack to find user-space filename/line number, and the procfunc function name
+    """
+
+    procfunc_frame = inspect.currentframe().f_back  # type: ignore
+    while (
+        module_path() in Path(procfunc_frame.f_back.f_code.co_filename).parents  # type: ignore
+    ):
+        procfunc_frame = procfunc_frame.f_back  # type: ignore
+    return (
+        procfunc_frame.f_back.f_code.co_filename,  # type: ignore
+        procfunc_frame.f_back.f_lineno,  # type: ignore
+        procfunc_frame.f_code.co_name,  # type: ignore
+    )
+
+
+def _has_unpreprocessed_inputs(node: cg.Node) -> bool:
+    any_args = any(isinstance(a, ProcNode) for a in node.args)
+    any_kwargs = any(isinstance(v, ProcNode) for v in node.kwargs.values())
+    return any_args or any_kwargs
+
+
+class ProcNode(Generic[T]):
+    """
+    Result datatype for all functions that return shader nodes, geometry nodes or compositor nodes.
+
+    ProcNode stores the data necessary to construct a blender nodegroup upon later execution.
+
+    ProcNode defines dunders to allow concise construction of nodegraphs e.g. __getattr__ and __add__, which map to appropriate blender nodes.
+    """
+
+    def __init__(
+        self,
+        node: cg.Node,
+        known_value_type: type | None = None,
+    ):
+        self._node = node
+
+        if known_value_type is not None:
+            logger.debug(f"{self} using provided known_value_type={known_value_type}")
+            self._node.metadata["known_value_type"] = known_value_type
+
+        if _has_unpreprocessed_inputs(node):
+            raise ValueError(
+                f"{node=} has inputs which are ProcNode, "
+                f"these should have been unwrapped to cg.Node {node.args} {node.kwargs}"
+            )
+
+        self._node.metadata["definition"] = _node_definition_metadata()
+
+    def astype(self, dtype: type) -> "ProcNode":
+        """
+        Marks a node as having been converted to a different internal data type, similarly to np.astype
+
+        Currently this just adds runtime NodeType data to help subsequent type-inferred functions/operators
+        make a correct choice of data_type.
+
+        e.g noise.color + (0.5, 0.5, 0.5) fails but noise.color.astype(t.Vector) + (0.5, 0.5, 0.5) works,
+        because `+` is defined for Vector but not Color
+        """
+
+        node = copy.copy(self._node)
+        node.metadata = copy.copy(self._node.metadata)
+        node.metadata["known_value_type"] = dtype
+
+        logger.debug(f"{self}.astype() using provided known_value_type={dtype}")
+        return ProcNode(node)
+
+    def __repr__(self):
+        # NOTE: dont change this to be anything verbose, it may slow down system
+        #   due to generating strings for debug logs (even if they arent actually printed)
+        return f"ProcNode({self.item()!r})"
+
+    @classmethod
+    def from_nodetype(
+        cls,
+        node_type: str,
+        inputs: dict[str, Any],
+        attrs: dict[str, Any],
+    ) -> "ProcNode":
+        def _unwrap(v: Any) -> cg.Node:
+            if isinstance(v, ProcNode):
+                return object.__getattribute__(v, "_node")
+            return v
+
+        inputs = pytree.PyTree(inputs).map(_unwrap).obj()
+
+        if any(isinstance(v, ProcNode) for v in attrs.values()):
+            raise ValueError(
+                f"Attrs {attrs} contains ProcNode, which is not allowed. Must specify a constant."
+            )
+
+        node = cg.ProceduralNode(node_type=node_type, attrs=attrs, kwargs=inputs)
+        node.metadata["definition"] = _node_definition_metadata()
+
+        return cls(node=node)
+
+    def item(self) -> cg.Node:
+        return object.__getattribute__(self, "_node")
+
+    def __post_init__(self):
+        for k, v in self.attrs.items():
+            if isinstance(v, (bpy.types.NodeInternal, ProcNode)):
+                raise ValueError(
+                    f"Node {self.type} has a {k} attribute that is a Node, which is not allowed. Must specify a constant."
+                )
+
+    def _output_socket(self, name: str) -> "ProcNode":
+        node = cg.GetAttributeNode(source=self.item(), attribute_name=name)
+        return ProcNode(node)
+
+    def _procnode_operator(
+        self,
+        op: OperatorType,
+        lhs: "ProcNode[T]",
+        rhs: "ProcNode[T] | T",
+        reverse: bool = False,
+    ) -> "ProcNode[T]":
+        rhs_unwrap = rhs.item() if isinstance(rhs, ProcNode) else rhs
+
+        if reverse:
+            args = (rhs_unwrap, lhs.item())
+        else:
+            args = (lhs.item(), rhs_unwrap)
+
+        node = cg.FunctionCallNode(
+            func=OPERATORS_TO_FUNCTIONS[op],
+            args=args,
+            kwargs={},
+            metadata=None,
+        )
+        return ProcNode(node)
+
+    def _getattr_xyz(
+        self: "ProcNode[pt.Vector]",
+        name: str,
+    ) -> "ProcNode[float]":
+        sep = ProcNode.from_nodetype(
+            node_type="ShaderNodeSeparateXYZ",
+            inputs={"Vector": self},
+            attrs={},
+        )
+        return sep._output_socket(name)
+
+    @property
+    def x(self: "ProcNode[pt.Vector]") -> "ProcNode[float]":
+        return self._getattr_xyz(name="x")
+
+    @property
+    def y(self: "ProcNode[pt.Vector]") -> "ProcNode[float]":
+        return self._getattr_xyz(name="y")
+
+    @property
+    def z(self: "ProcNode[pt.Vector]") -> "ProcNode[float]":
+        return self._getattr_xyz(name="z")
+
+    def __add__(self, other: "ProcNode[T] | T | tuple") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.ADD, self, other)
+
+    def __radd__(self, other: "ProcNode[T] | T | tuple") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.ADD, self, other, reverse=True)
+
+    def __sub__(self, other: "ProcNode[T] | T | tuple") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.SUB, self, other)
+
+    def __rsub__(self, other: "ProcNode[T] | T | tuple") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.SUB, self, other, reverse=True)
+
+    def __mul__(self, other: "ProcNode[T] | T | tuple") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.MUL, self, other)
+
+    def __rmul__(self, other: "ProcNode[T] | T | tuple") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.MUL, self, other, reverse=True)
+
+    def __truediv__(self, other: "ProcNode[T] | T | tuple") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.DIV, self, other)
+
+    def __rtruediv__(self, other: "ProcNode[T] | T | tuple") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.DIV, self, other, reverse=True)
+
+    def __pow__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.POW, self, other)
+
+    def __rpow__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.POW, self, other, reverse=True)
+
+    def __mod__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.MOD, self, other)
+
+    def __rmod__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.MOD, self, other, reverse=True)
+
+    def __lt__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.LESS_THAN, self, other)
+
+    def __rlt__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(
+            OperatorType.LESS_THAN, self, other, reverse=True
+        )
+
+    def __le__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.LESS_THAN_EQUAL, self, other)
+
+    def __rle__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(
+            OperatorType.LESS_THAN_EQUAL, self, other, reverse=True
+        )
+
+    def __gt__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.GREATER_THAN, self, other)
+
+    def __rgt__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(
+            OperatorType.GREATER_THAN, self, other, reverse=True
+        )
+
+    def __ge__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.GREATER_THAN_EQUAL, self, other)
+
+    def __rge__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(
+            OperatorType.GREATER_THAN_EQUAL, self, other, reverse=True
+        )
+
+    def __eq__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.EQUAL, self, other)
+
+    def __req__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.EQUAL, self, other, reverse=True)
+
+    def __ne__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(OperatorType.NOT_EQUAL, self, other)
+
+    def __rne__(self, other: "ProcNode[T] | T") -> "ProcNode[T]":
+        return self._procnode_operator(
+            OperatorType.NOT_EQUAL, self, other, reverse=True
+        )
+
+
+def node_definition_context_message(node: cg.Node):
+    metadata = object.__getattribute__(node, "metadata")
+    lineno_metadata = metadata.get("definition", None)
+    if lineno_metadata is None:
+        return ""
+    file, lineno, procfunc_name = lineno_metadata
+    return f" {procfunc_name}() call on {file}:{lineno} "
+
+
+TSocketVal = TypeVar("TSocketVal")
+SocketOrVal = Union[ProcNode[TSocketVal], TSocketVal]
+
+
+class Instances:
+    pass
+
+
+Points = Union[pt.MeshObject, pt.CurveObject]
+
+Geometry = Union[pt.MeshObject, pt.CurveObject, Instances, pt.VolumeObject]
+
+AnyShaderDataVal = Union[
+    pt.Vector,
+    pt.Color,
+    float,
+]
+
+AnyDataVal = Union[AnyShaderDataVal, int, str, bool, pt.Matrix, pt.Quaternion]
+"""
+Union of all types that are data-like in a geometrynodes context
+IE pretty much all the geonodes sockettypes except specialones like Material or Object
+"""
+
+AnyAssetVal = Union[
+    pt.Object,
+    pt.Collection,
+    pt.Material,
+    pt.Texture,
+]
+"""
+Union of all types that are object-like in a geometrynodes context
+"""
+
+AnyVal = Union[AnyDataVal, AnyAssetVal]
+
+
+class Shader:
+    """
+    Used only for type-annotating nodes as returning a shader.
+
+    Anythnig that would be a green socket in a SHADER nodegraph should be ProcNode[Shader]
+    """
+
+    pass
+
+
+__all__ = [
+    "ProcNode",
+    "Shader",
+    "SocketOrVal",
+    "AnyShaderDataVal",
+    "AnyDataVal",
+    "AnyAssetVal",
+    "node_definition_context_message",
+]

procfunc/ops/__init__.py

@@ -0,0 +1,35 @@
+from pandas import read_json as _read_json
+
+from procfunc.tracer import autowrap_module as _autowrap
+from procfunc.util.manifest import module_path
+
+from . import (
+    attr,
+    collection,
+    curve,
+    file,
+    mesh,
+    modifier,
+    object,
+    primitives,
+    uv,
+)
+
+_autowrap(collection)
+_autowrap(attr)
+_autowrap(curve)
+_autowrap(primitives)
+
+OPS_MANIFEST_PATH = module_path() / "ops" / "manifest.json"
+OPS_MANIFEST = _read_json(OPS_MANIFEST_PATH)
+
+__all__ = [
+    "collection",
+    "file",
+    "mesh",
+    "modifier",
+    "object",
+    "attr",
+    "curve",
+    "uv",
+]

procfunc/ops/_util.py

@@ -0,0 +1,275 @@
+import logging
+from typing import Any, Callable, Literal, TypeVar
+
+import bpy
+import numpy as np
+
+from procfunc import types as t
+from procfunc.util.bpy_info import bpy_nocollide_data_name
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_objs(
+    objs: t.Object | list[t.Object] | None, active: t.Object | None
+) -> tuple[list[bpy.types.Object], bpy.types.Object]:
+    if objs is None:
+        assert active is not None
+        objs = [active.item()]
+    elif isinstance(objs, t.Object):
+        objs = [objs.item()]
+    elif isinstance(objs, list):
+        objs = [obj.item() for obj in objs]
+    else:
+        raise TypeError(f"Unexpected type for objs: {type(objs)}")
+
+    if active is None:
+        assert len(objs) > 0
+        active = objs[0]
+    elif isinstance(active, t.Asset):
+        active = active.item()
+
+    if active not in objs:
+        objs.append(active)
+
+    assert active is not None
+    assert len(objs) > 0
+    assert not isinstance(active, t.Asset)
+
+    return objs, active
+
+
+def _assert_op_finished(result: Any, name: str):
+    if (
+        not isinstance(result, set)
+        or len(result) != 1
+        or next(iter(result)) != "FINISHED"
+    ):
+        raise ValueError(f"{name} got status {result}, expected {'FINISHED'}")
+
+
+def execute_object_op(
+    operator: Callable,
+    objs: t.Object | list[t.Object] | None = None,
+    active: t.Object | None = None,
+    description: str = "",
+    **kwargs: Any,
+):
+    assert not isinstance(active, bpy.types.Object)
+
+    kwargs = {k: v.item() if isinstance(v, t.Asset) else v for k, v in kwargs.items()}
+
+    objs, active = _parse_objs(objs, active)
+
+    # TODO: more efficient via temp_override?
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.context.view_layer.objects.active = active
+    bpy.ops.object.mode_set(mode="OBJECT")
+    for o in objs:
+        o.select_set(True)
+
+    result = operator(**kwargs)
+    _assert_op_finished(result, description)
+
+
+def _apply_selection_masks(
+    obj: bpy.types.Object,
+    vertex_mask: np.ndarray | None = None,
+    edge_mask: np.ndarray | None = None,
+    face_mask: np.ndarray | None = None,
+):
+    """
+    Apply one selection mask of some domain to the object
+
+    If no masks are provided, all elements of the domain are selected.
+
+    NOTE: function must start AND end in EDIT mode
+    """
+
+    n_masks = sum(mask is not None for mask in [vertex_mask, edge_mask, face_mask])
+
+    if n_masks == 0:
+        logger.debug(f"No edit-masks provided for {obj.name}, selecting all")
+        bpy.ops.mesh.select_all(action="SELECT")
+        return
+    if n_masks > 1:
+        raise ValueError(
+            "Only one of vertex_mask, edge_mask, or face_mask can be provided"
+        )
+
+    bpy.ops.mesh.select_all(action="DESELECT")
+
+    # must select the type of mask we're applying, otherwise re-entering editmode will incorrectly
+    # convert face/edge masks into vertex masks, which is lossy.
+    select_type = (
+        "VERT"
+        if vertex_mask is not None
+        else "EDGE"
+        if edge_mask is not None
+        else "FACE"
+    )
+    bpy.ops.mesh.select_mode(type=select_type)
+
+    # foreach_set operations must run in object mode then switch back to edit mode to re-sync them
+    bpy.ops.object.mode_set(mode="OBJECT")
+
+    if vertex_mask is not None:
+        assert vertex_mask.shape == (len(obj.data.vertices),)
+        assert vertex_mask.dtype == bool
+        obj.data.vertices.foreach_set("select", vertex_mask)
+
+    if edge_mask is not None:
+        assert edge_mask.shape == (len(obj.data.edges),)
+        assert edge_mask.dtype == bool
+        obj.data.edges.foreach_set("select", edge_mask)
+
+    if face_mask is not None:
+        assert face_mask.shape == (len(obj.data.polygons),)
+        assert face_mask.dtype == bool
+        obj.data.polygons.foreach_set("select", face_mask)
+
+    bpy.ops.object.mode_set(mode="EDIT")
+
+
+def extract_face_mask(
+    obj: t.MeshObject,
+) -> np.ndarray:
+    obj = obj.item()
+    face_mask = np.zeros(len(obj.data.polygons), dtype=bool)
+    obj.data.polygons.foreach_get("select", face_mask)
+    return face_mask
+
+
+def extract_edge_mask(
+    obj: t.MeshObject,
+) -> np.ndarray:
+    obj = obj.item()
+    edge_mask = np.zeros(len(obj.data.edges), dtype=bool)
+    obj.data.edges.foreach_get("select", edge_mask)
+    return edge_mask
+
+
+def extract_vertex_mask(
+    obj: t.MeshObject,
+) -> np.ndarray:
+    obj = obj.item()
+    vertex_mask = np.zeros(len(obj.data.vertices), dtype=bool)
+    obj.data.vertices.foreach_get("select", vertex_mask)
+    return vertex_mask
+
+
+def execute_mesh_op(
+    operator: Callable,
+    obj: t.MeshObject,
+    vertex_mask: np.ndarray | None = None,
+    edge_mask: np.ndarray | None = None,
+    face_mask: np.ndarray | None = None,
+    empty_mask_mode: Literal["return", "execute", "error"] = "return",
+    description: str = "",
+    **kwargs: Any,
+):
+    obj = obj.item()
+
+    logger.debug(f"Executing {operator} on {obj.name=}")
+
+    empty_mask = next(
+        (
+            mask
+            for mask in [vertex_mask, edge_mask, face_mask]
+            if mask is not None and mask.sum() == 0
+        ),
+        None,
+    )
+    if empty_mask is not None and empty_mask_mode == "error":
+        raise ValueError(
+            f"Empty mask provided for {operator} on {obj.name=} with {empty_mask_mode=}"
+        )
+
+    # TODO: would rather do this with context overrides,
+    #   but it messes up the active_object for cases like ops.mesh.separate()
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.context.view_layer.objects.active = obj
+    bpy.ops.object.mode_set(mode="OBJECT")
+    obj.select_set(True)
+
+    bpy.ops.object.mode_set(mode="EDIT")
+    _apply_selection_masks(obj, vertex_mask, edge_mask, face_mask)
+
+    try:
+        if empty_mask is not None:
+            if empty_mask_mode == "return":
+                return
+            elif empty_mask_mode == "execute":
+                pass
+            elif empty_mask_mode == "error":
+                raise ValueError(
+                    f"Empty mask provided for {operator} on {obj.name=} with {empty_mask_mode=}"
+                )
+            else:
+                raise ValueError(f"Invalid empty_mask_mode: {empty_mask_mode}")
+
+        result = operator(**kwargs)
+    finally:
+        bpy.ops.object.mode_set(mode="OBJECT")
+
+    _assert_op_finished(result, description)
+
+    return result
+
+
+TModifyObject = TypeVar("TModifyObject", t.MeshObject, t.CurveObject)
+
+
+def modify(
+    obj: TModifyObject,
+    modifier_type: str,
+    setitem_keyvals: dict[str, Any] | None = None,
+    _skip_apply: bool = False,
+    **setattr_keyvals: Any,
+) -> TModifyObject:
+    if setitem_keyvals is None:
+        setitem_keyvals = {}
+
+    mod_name = bpy_nocollide_data_name(modifier_type, obj.item().modifiers)
+    modifier = obj.item().modifiers.new(mod_name, modifier_type)
+    if modifier is None:
+        raise ValueError(
+            "modifier.new() returned None, blender might not allow "
+            f"{modifier_type=} for {obj.item().type=}"
+        )
+    modifier.show_viewport = False
+
+    for key, value in setitem_keyvals.items():
+        assert not isinstance(value, t.Asset), (
+            "TODO may need to conver Assets to item for modifier[x] = y"
+        )
+        modifier[key] = value
+
+    objs = []
+    for key, value in setattr_keyvals.items():
+        if isinstance(value, t.Object):
+            objs.append(value)
+        if isinstance(value, t.Asset):
+            value = value.item()
+
+        if hasattr(modifier, key):
+            setattr(modifier, key, value)
+        else:
+            raise AttributeError(f"Modifier {modifier_type} has no attribute {key}")
+
+    if _skip_apply:
+        return obj
+
+    execute_object_op(
+        bpy.ops.object.modifier_apply,
+        objs=objs,
+        active=obj,
+        modifier=modifier.name,
+    )
+
+    if mod_name in obj.item().modifiers:
+        raise ValueError(
+            f"Modifier {modifier_type} failed to execute with {setitem_keyvals=} and {setattr_keyvals=}"
+        )
+
+    return obj