nodebpy 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

nodebpy/builder.py

@@ -1,10 +1,6 @@
 from __future__ import annotations
 
-from ast import Return
-from typing import TYPE_CHECKING, Any, ClassVar, Literal
-
-from arrangebpy.arrange.routing import Socket
-from numpy import isin
+from typing import TYPE_CHECKING, Any, ClassVar, Iterable, Literal
 
 if TYPE_CHECKING:
     from .nodes import Math, VectorMath
@@ -25,7 +21,6 @@ from .types import (
     TYPE_INPUT_ALL,
     FloatInterfaceSubtypes,
     IntegerInterfaceSubtypes,
-    NodeBooleanMathItems,
     StringInterfaceSubtypes,
     VectorInterfaceSubtypes,
     _AttributeDomains,
@@ -121,8 +116,9 @@ class OutputInterfaceContext(DirectionalContext):
 class TreeBuilder:
     """Builder for creating Blender geometry node trees with a clean Python API."""
 
-    _active_tree: ClassVar["TreeBuilder | None"] = None
-    _previous_tree: ClassVar["TreeBuilder | None"] = None
+    _tree_contexts: ClassVar["list[TreeBuilder]"] = []
+    # _active_tree: ClassVar["TreeBuilder | None"] = None
+    # _previous_tree: ClassVar["list[TreeBuilder]"] = list()
     just_added: "Node | None" = None
 
     def __init__(
@@ -139,16 +135,22 @@ class TreeBuilder:
         self.outputs = OutputInterfaceContext(self)
         self._arrange = arrange
 
+    def activate_tree(self) -> None:
+        """Make this tree the active tree for all new node creation."""
+        TreeBuilder._tree_contexts.append(self)
+
+    def deactivate_tree(self) -> None:
+        """Whatever tree ws previously active is set to be the active one (or None if no previously active tree)."""
+        TreeBuilder._tree_contexts.pop()
+
     def __enter__(self):
-        TreeBuilder._previous_tree = TreeBuilder._active_tree
-        TreeBuilder._active_tree = self
+        self.activate_tree()
         return self
 
     def __exit__(self, *args):
         if self._arrange:
             self.arrange()
-        TreeBuilder._active_tree = TreeBuilder._previous_tree
-        TreeBuilder._previous_tree = None
+        self.deactivate_tree()
 
     @property
     def nodes(self) -> Nodes:
@@ -244,7 +246,9 @@ class NodeBuilder:
 
     def __init__(self):
         # Get active tree from context manager
-        tree = TreeBuilder._active_tree
+        tree = (
+            TreeBuilder._tree_contexts[-1] if len(TreeBuilder._tree_contexts) else None
+        )
         if tree is None:
             raise RuntimeError(
                 f"Node '{self.__class__.__name__}' must be created within a TreeBuilder context manager.\n"
@@ -344,18 +348,20 @@ class NodeBuilder:
 
         raise SocketError("No compatible output sockets found")
 
-    @staticmethod
     def _find_best_socket_pair(
-        source: "NodeBuilder | NodeSocket", target: "NodeBuilder | NodeSocket"
+        self,
+        source: "NodeBuilder | SocketLinker | NodeSocket",
+        target: "NodeBuilder | SocketLinker | NodeSocket",
     ) -> tuple[NodeSocket, NodeSocket]:
         """Find the best possible compatible pair of sockets between two nodes, looking only at the
         the currently available outputs from the source and the inputs from the target"""
         possible_combos = []
-        if isinstance(source, NodeBuilder):
+        if isinstance(source, (NodeBuilder, SocketLinker)):
             outputs = source._available_outputs
-        else:
+        elif isinstance(source, NodeSocket):
             outputs = [source]
-        if isinstance(target, NodeBuilder):
+
+        if isinstance(target, (NodeBuilder, SocketLinker)):
             inputs = target._available_inputs
         else:
             inputs = [target]
@@ -596,7 +602,9 @@ class DynamicInputsMixin:
     _socket_data_types: tuple[str]
     _type_map: dict[str, str] = {}
 
-    def _match_compatible_data(self, *sockets: NodeSocket) -> tuple[NodeSocket, str]:
+    def _match_compatible_data(
+        self, sockets: Iterable[NodeSocket]
+    ) -> tuple[NodeSocket, str]:
         possible = []
         for socket in sockets:
             compatible = SOCKET_COMPATIBILITY.get(socket.type, ())
@@ -612,8 +620,8 @@ class DynamicInputsMixin:
         raise SocketError("No compatible socket found")
 
     def _find_best_socket_pair(
-        self, source: NodeBuilder, target: NodeBuilder
-    ) -> tuple[NodeSocket, NodeSocket] | None:
+        self, source: NodeBuilder | NodeSocket, target: NodeBuilder | NodeSocket
+    ) -> tuple[NodeSocket, NodeSocket]:
         try:
             return super()._find_best_socket_pair(source, target)
         except SocketError:
@@ -629,9 +637,9 @@ class DynamicInputsMixin:
                     target.inputs[target_name].socket,
                 )
 
-            for target_name, source_socket in new_sockets.items():
-                target_socket = target.inputs[target_name].socket
-                return (source_socket, target_socket)
+            # for target_name, source_socket in new_sockets.items():
+            #     target_socket = target.inputs[target_name].socket
+            #     return (source_socket, target_socket)
 
     # def _best_output_socket(self, type: str) -> NodeSocket:
     #     # compatible = SOCKET_COMPATIBILITY.get(type, ())
@@ -655,7 +663,7 @@ class DynamicInputsMixin:
                 items[arg._default_output_socket.name] = arg
         items.update(kwargs)
         for key, source in items.items():
-            socket_source, type = self._match_compatible_data(*source.node.outputs)
+            socket_source, type = self._match_compatible_data(source._available_outputs)
             if type in self._type_map:
                 type = self._type_map[type]
             socket = self._add_socket(name=key, type=type)
@@ -676,6 +684,10 @@ class SocketLinker(NodeBuilder):
     def _available_outputs(self) -> list[NodeSocket]:
         return [self.socket]
 
+    @property
+    def _available_inputs(self) -> list[NodeSocket]:
+        return [self.socket]
+
     @property
     def type(self) -> SOCKET_TYPES:
         return self.socket.type  # type: ignore

nodebpy/nodes/__init__.py

@@ -18,6 +18,9 @@ from .manual import (
     SimulationInput,
     SimulationOutput,
     SimulationZone,
+    ForEachGeometryElementInput,
+    ForEachGeometryElementOutput,
+    ForEachGeometryElementZone,
     FormatString,
     Value,
     AccumulateField,
@@ -414,6 +417,9 @@ __all__ = (
     "FlipFaces",
     "FloatCurve",
     "FloatToInteger",
+    "ForEachGeometryElementInput",
+    "ForEachGeometryElementOutput",
+    "ForEachGeometryElementZone",
     "FormatString",
     "GaborTexture",
     "Gamma",

nodebpy/nodes/attribute.py

@@ -11,13 +11,16 @@ from ..types import (
     TYPE_INPUT_STRING,
     TYPE_INPUT_ROTATION,
     TYPE_INPUT_COLOR,
+    TYPE_INPUT_MATRIX,
     TYPE_INPUT_VALUE,
     TYPE_INPUT_VECTOR,
 )
 
 
 class BlurAttribute(NodeBuilder):
-    """Mix attribute values of neighboring elements"""
+    """
+    Mix attribute values of neighboring elements
+    """
 
     _bl_idname = "GeometryNodeBlurAttribute"
     node: bpy.types.GeometryNodeBlurAttribute
@@ -109,7 +112,9 @@ class BlurAttribute(NodeBuilder):
 
 
 class DomainSize(NodeBuilder):
-    """Retrieve the number of elements in a geometry for each attribute domain"""
+    """
+    Retrieve the number of elements in a geometry for each attribute domain
+    """
 
     _bl_idname = "GeometryNodeAttributeDomainSize"
     node: bpy.types.GeometryNodeAttributeDomainSize
@@ -181,7 +186,9 @@ class DomainSize(NodeBuilder):
 
 
 class RemoveNamedAttribute(NodeBuilder):
-    """Delete an attribute with a specified name from a geometry. Typically used to optimize performance"""
+    """
+    Delete an attribute with a specified name from a geometry. Typically used to optimize performance
+    """
 
     _bl_idname = "GeometryNodeRemoveAttribute"
     node: bpy.types.GeometryNodeRemoveAttribute
@@ -219,7 +226,9 @@ class RemoveNamedAttribute(NodeBuilder):
 
 
 class StoreNamedAttribute(NodeBuilder):
-    """Store the result of a field on a geometry as an attribute with the specified name"""
+    """
+    Store the result of a field on a geometry as an attribute with the specified name
+    """
 
     _bl_idname = "GeometryNodeStoreNamedAttribute"
     node: bpy.types.GeometryNodeStoreNamedAttribute
@@ -360,6 +369,74 @@ class StoreNamedAttribute(NodeBuilder):
             value=value,
         )
 
+    @classmethod
+    def matrix(
+        cls,
+        geometry: TYPE_INPUT_GEOMETRY = None,
+        selection: TYPE_INPUT_BOOLEAN = True,
+        name: TYPE_INPUT_STRING = "",
+        value: TYPE_INPUT_MATRIX = None,
+    ) -> "StoreNamedAttribute":
+        """Create Store Named Attribute with operation '4x4 Matrix'."""
+        return cls(
+            data_type="FLOAT4X4",
+            geometry=geometry,
+            selection=selection,
+            name=name,
+            value=value,
+        )
+
+    @classmethod
+    def int8(
+        cls,
+        geometry: TYPE_INPUT_GEOMETRY = None,
+        selection: TYPE_INPUT_BOOLEAN = True,
+        name: TYPE_INPUT_STRING = "",
+        value: TYPE_INPUT_INT = 0,
+    ) -> "StoreNamedAttribute":
+        """Create Store Named Attribute with operation '8-Bit Integer'."""
+        return cls(
+            data_type="INT8",
+            geometry=geometry,
+            selection=selection,
+            name=name,
+            value=value,
+        )
+
+    @classmethod
+    def vector2(
+        cls,
+        geometry: TYPE_INPUT_GEOMETRY = None,
+        selection: TYPE_INPUT_BOOLEAN = True,
+        name: TYPE_INPUT_STRING = "",
+        value: TYPE_INPUT_VECTOR = None,
+    ) -> "StoreNamedAttribute":
+        """Create Store Named Attribute with operation '2D Vector'."""
+        return cls(
+            data_type="FLOAT2",
+            geometry=geometry,
+            selection=selection,
+            name=name,
+            value=value,
+        )
+
+    @classmethod
+    def byte_color(
+        cls,
+        geometry: TYPE_INPUT_GEOMETRY = None,
+        selection: TYPE_INPUT_BOOLEAN = True,
+        name: TYPE_INPUT_STRING = "",
+        value: TYPE_INPUT_COLOR = None,
+    ) -> "StoreNamedAttribute":
+        """Create Store Named Attribute with operation 'Byte Color'."""
+        return cls(
+            data_type="BYTE_COLOR",
+            geometry=geometry,
+            selection=selection,
+            name=name,
+            value=value,
+        )
+
     @classmethod
     def point(
         cls,
@@ -411,6 +488,23 @@ class StoreNamedAttribute(NodeBuilder):
             value=value,
         )
 
+    @classmethod
+    def face_corner(
+        cls,
+        geometry: TYPE_INPUT_GEOMETRY = None,
+        selection: TYPE_INPUT_BOOLEAN = True,
+        name: TYPE_INPUT_STRING = "",
+        value: TYPE_INPUT_COLOR = None,
+    ) -> "StoreNamedAttribute":
+        """Create Store Named Attribute with operation 'Face Corner'."""
+        return cls(
+            domain="CORNER",
+            geometry=geometry,
+            selection=selection,
+            name=name,
+            value=value,
+        )
+
     @classmethod
     def spline(
         cls,

nodebpy/nodes/color.py

@@ -9,7 +9,9 @@ from ..types import (
 
 
 class Gamma(NodeBuilder):
-    """Apply a gamma correction"""
+    """
+    Apply a gamma correction
+    """
 
     _bl_idname = "ShaderNodeGamma"
     node: bpy.types.ShaderNodeGamma
@@ -41,7 +43,9 @@ class Gamma(NodeBuilder):
 
 
 class RgbCurves(NodeBuilder):
-    """Apply color corrections for each color channel"""
+    """
+    Apply color corrections for each color channel
+    """
 
     _bl_idname = "ShaderNodeRGBCurve"
     node: bpy.types.ShaderNodeRGBCurve