Add built-in node arrange and documentation

Author: carson-katri

api/arrange.py

@@ -0,0 +1,61 @@
+import bpy
+import typing
+
+def _arrange(node_tree, padding: typing.Tuple[float, float] = (50, 25)):
+    # Organize the nodes into columns based on their links.
+    columns: typing.List[typing.List[typing.Any]] = []
+    def contains_link(node, column):
+        return any(
+            any(
+                any(link.from_node == node for link in input.links)
+                for input in n.inputs
+            )
+            for n in column
+        )
+    for node in reversed(node_tree.nodes):
+        if (x := next(
+            filter(
+                lambda x: contains_link(node, x[1]),
+                enumerate(columns)
+            ),
+            None
+        )) is not None:
+            if x[0] > 0:
+                columns[x[0] - 1].append(node)
+            else:
+                columns.insert(x[0], [node])
+        else:
+            if len(columns) == 0:
+                columns.append([node])
+            else:
+                columns[len(columns) - 1].append(node)
+    
+    # Arrange the columns, computing the size of the node manually so arrangement can be done without UI being visible.
+    UI_SCALE = bpy.context.preferences.view.ui_scale
+    NODE_HEADER_HEIGHT = 20
+    NODE_LINK_HEIGHT = 28
+    NODE_PROPERTY_HEIGHT = 28
+    NODE_VECTOR_HEIGHT = 84
+    x = 0
+    for col in columns:
+        largest_width = 0
+        y = 0
+        for node in col:
+            node.update()
+            input_count = len(list(filter(lambda i: i.enabled, node.inputs)))
+            output_count = len(list(filter(lambda i: i.enabled, node.outputs)))
+            parent_props = [prop.identifier for base in type(node).__bases__ for prop in base.bl_rna.properties]
+            properties_count = len([prop for prop in node.bl_rna.properties if prop.identifier not in parent_props])
+            unset_vector_count = len(list(filter(lambda i: i.enabled and i.type == 'VECTOR' and len(i.links) == 0, node.inputs)))
+            node_height = (
+                NODE_HEADER_HEIGHT \
+                + (output_count * NODE_LINK_HEIGHT) \
+                + (properties_count * NODE_PROPERTY_HEIGHT) \
+                + (input_count * NODE_LINK_HEIGHT) \
+                + (unset_vector_count * NODE_VECTOR_HEIGHT)
+            ) * UI_SCALE
+            if node.width > largest_width:
+                largest_width = node.width
+            node.location = (x, y)
+            y -= node_height + padding[1]
+        x += largest_width + padding[0]

api/node_mapper.py

@@ -15,10 +15,12 @@ class OutputsList(dict):
 
 def build_node(node_type):
     def build(_primary_arg=None, **kwargs):
-        for k, v in kwargs.items():
+        for k, v in kwargs.copy().items():
             if isinstance(v, InputGroup):
                 kwargs = { **kwargs, **v.__dict__ }
                 del kwargs[k]
+            if v is None:
+                del kwargs[k]
         node = State.current_node_tree.nodes.new(node_type.__name__)
         if _primary_arg is not None:
             State.current_node_tree.links.new(_primary_arg._socket, node.inputs[0])
@@ -316,6 +318,7 @@ def tree(builder):
   Marks a function as a node tree.
   \"\"\"
   pass
+_SomeType = TypeVar('_SomeType', bound='Type')
 class Type:
   def __add__(self, other) -> Type: return self
   def __radd__(self, other) -> Type: return self
@@ -333,6 +336,11 @@ def __lt__(self, other) -> Type: return self
   def __le__(self, other) -> Type: return self
   def __gt__(self, other) -> Type: return self
   def __ge__(self, other) -> Type: return self
+  def __invert__(self) -> Type: return self
+  def __getitem__(
+    self,
+    subscript: _SomeType | slice | Tuple[_SomeType | slice, SampleMode]
+  ) -> Type: return self
   x = Type()
   y = Type()
   z = Type()

api/static/sample_mode.py

@@ -0,0 +1,6 @@
+import enum
+
+class SampleMode(enum.IntEnum):
+    INDEX = 0
+    NEAREST_SURFACE = 1
+    NEAREST = 2

api/tree.py

@@ -10,8 +10,12 @@
 from .static.attribute import *
 from .static.expression import *
 from .static.input_group import *
+from .static.sample_mode import *
+from .arrange import _arrange
 
 def _as_iterable(x):
+    if isinstance(x, Type):
+        return [x,]
     try:
         return iter(x)
     except TypeError:
@@ -100,35 +104,7 @@ def validate_param(param):
             node_group.outputs.new(result.socket_type, 'Result')
             link = node_group.links.new(result._socket, group_output_node.inputs[i])
 
-        # Attempt to run the "Node Arrange" add-on on the tree.
-        try:
-            for area in bpy.context.screen.areas:
-                for space in area.spaces:
-                    if space.type == 'NODE_EDITOR':
-                        space.node_tree = node_group
-                        with bpy.context.temp_override(area=area, space=space, space_data=space):
-                            ntree = node_group
-                            ntree.nodes[0].select = True
-                            ntree.nodes.active = ntree.nodes[0]
-                            n_groups = []
-                            for i in ntree.nodes:
-                                if i.type == 'GROUP':
-                                    n_groups.append(i)
-
-                            while n_groups:
-                                j = n_groups.pop(0)
-                                node_arrange.nodes_iterate(j.node_tree)
-                                for i in j.node_tree.nodes:
-                                    if i.type == 'GROUP':
-                                        n_groups.append(i)
-
-                            node_arrange.nodes_iterate(ntree)
-
-                            # arrange nodes + this center nodes together
-                            if bpy.context.scene.node_center:
-                                node_arrange.nodes_center(ntree)
-        except:
-            pass
+        _arrange(node_group)
 
         # Return a function that creates a NodeGroup node in the tree.
         # This lets @trees be used in other @trees via simple function calls.

api/types.py

@@ -1,7 +1,9 @@
 import bpy
 from bpy.types import NodeSocketStandard
 import nodeitems_utils
+import enum
 from .state import State
+from .static.sample_mode import SampleMode
 import geometry_script
 
 def map_case_name(i):
@@ -178,6 +180,53 @@ def capture(self, value, **kwargs):
     def transfer(self, attribute, **kwargs):
         data_type = socket_type_to_data_type(attribute._socket.type)
         return self.transfer_attribute(data_type=data_type, attribute=attribute, **kwargs)
+    
+    def __getitem__(self, subscript):
+        if isinstance(subscript, tuple):
+            accessor = subscript[0]
+            args = subscript[1:]
+        else:
+            accessor = subscript
+            args = []
+        sample_mode = SampleMode.INDEX if len(args) < 1 else args[0]
+        domain = 'POINT' if len(args) < 2 else (args[1].value if isinstance(args[1], enum.Enum) else args[1])
+        sample_position = None
+        sampling_index = None
+        if isinstance(accessor, slice):
+            data_type = socket_type_to_data_type(accessor.start._socket.type)
+            value = accessor.start
+            match sample_mode:
+                case SampleMode.INDEX:
+                    sampling_index = accessor.stop
+                case SampleMode.NEAREST_SURFACE:
+                    sample_position = accessor.stop
+                case SampleMode.NEAREST:
+                    sample_position = accessor.stop
+            if accessor.step is not None:
+                domain = accessor.step.value if isinstance(accessor.step, enum.Enum) else accessor.step
+        else:
+            data_type = socket_type_to_data_type(accessor._socket.type)
+            value = accessor
+        match sample_mode:
+            case SampleMode.INDEX:
+                return self.sample_index(
+                    data_type=data_type,
+                    domain=domain,
+                    value=value,
+                    index=sampling_index or geometry_script.index()
+                )
+            case SampleMode.NEAREST_SURFACE:
+                return self.sample_nearest_surface(
+                    data_type=data_type,
+                    value=value,
+                    sample_position=sample_position or geometry_script.position()
+                )
+            case SampleMode.NEAREST:
+                return self.sample_index(
+                    data_type=data_type,
+                    value=value,
+                    index=self.sample_nearest(domain=domain, sample_position=sample_position or geometry_script.position())
+                )
 
 for standard_socket in list(filter(lambda x: 'NodeSocket' in x, dir(bpy.types))):
     name = standard_socket.replace('NodeSocket', '')

book/src/api/advanced-scripting/attributes.md

@@ -1,6 +1,6 @@
 # Attributes
 
-An important concept in Geometry Nodes is attributes. Many trees capture attributes or transfer them from one geometry to another.
+An important concept in Geometry Nodes is attributes. Many trees capture attributes or transfer them from one domain to another.
 
 When using these methods, the `data_type` argument must be correctly specified for the transfer to work as intended.
 
@@ -65,4 +65,25 @@ my_custom_attribute = Attribute(
 geometry = my_custom_attribute.store(geometry, 0.5)
 # Use the value by calling the attribute
 geometry = geometry.set_position(offset=my_custom_attribute())
-```
+```
+
+## Attribute Sampling
+In Blender 3.4+, transfer attribute was replaced with a few separate nodes: *Sample Index*, *Sample Nearest*, and *Sample Nearest Surface*.
+
+To avoid inputting data types and geometry manually, you can use the custom `Geometry` subscript.
+
+The structure for these subscripts is:
+
+```python
+geometry[value : index or sample position : domain, mode, domain]
+```
+
+Only the value argument is required. Other arguments can be supplied as needed.
+
+```python
+geometry[value]
+geometry[value : sample_position, SampleMode.NEAREST]
+geometry[value : index() + 1 : SampleIndex.Domain.EDGE]
+```
+
+Try passing different arguments and see how the resulting nodes are created.