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

nodebpy/builder.py

@@ -1,15 +1,10 @@
 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
 
-import arrangebpy
 import bpy
 from bpy.types import (
     GeometryNodeTree,
@@ -18,6 +13,7 @@ from bpy.types import (
     NodeSocket,
 )
 
+from .lib.nodearrange import arrange
 from .types import (
     LINKABLE,
     SOCKET_COMPATIBILITY,
@@ -25,7 +21,6 @@ from .types import (
     TYPE_INPUT_ALL,
     FloatInterfaceSubtypes,
     IntegerInterfaceSubtypes,
-    NodeBooleanMathItems,
     StringInterfaceSubtypes,
     VectorInterfaceSubtypes,
     _AttributeDomains,
@@ -121,12 +116,16 @@ 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]"] = []
     just_added: "Node | None" = None
+    collapse: bool = False
 
     def __init__(
-        self, tree: GeometryNodeTree | str = "Geometry Nodes", arrange: bool = True
+        self,
+        tree: GeometryNodeTree | str = "Geometry Nodes",
+        *,
+        collapse: bool = False,
+        arrange: bool = True,
     ):
         if isinstance(tree, str):
             self.tree = bpy.data.node_groups.new(tree, "GeometryNodeTree")
@@ -138,30 +137,35 @@ class TreeBuilder:
         self.inputs = InputInterfaceContext(self)
         self.outputs = OutputInterfaceContext(self)
         self._arrange = arrange
+        self.collapse = collapse
+
+    @property
+    def nodes(self) -> Nodes:
+        return self.tree.nodes
+
+    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
-
-    @property
-    def nodes(self) -> Nodes:
-        return self.tree.nodes
+        self.deactivate_tree()
 
     def __len__(self) -> int:
         return len(self.nodes)
 
     def arrange(self):
-        settings = arrangebpy.LayoutSettings(
-            horizontal_spacing=200, vertical_spacing=200, align_top_layer=True
-        )
-        arrangebpy.sugiyama_layout(self.tree, settings)
+        arrange.sugiyama.sugiyama_layout(self.tree)
+        arrange.sugiyama.config.reset()
 
     def _repr_markdown_(self) -> str | None:
         """
@@ -228,7 +232,9 @@ class TreeBuilder:
         return link
 
     def add(self, name: str) -> Node:
-        return self.tree.nodes.new(name)
+        node = self.tree.nodes.new(name)
+        node.hide = self.collapse
+        return node
 
 
 class NodeBuilder:
@@ -244,7 +250,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 +352,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 +606,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 +624,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 +641,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 +667,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 +688,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/lib/nodearrange/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: GPL-2.0-or-later
+from .arrange import sugiyama