Support prefixes for InputGroups

Author: carson-katri

api/node_mapper.py

@@ -2,6 +2,7 @@
 import bl_ui
 import itertools
 import enum
+import re
 from .state import State
 from .types import *
 from .static.input_group import InputGroup
@@ -299,6 +300,8 @@ def type_symbol(t):
         def enum_namespace(k):
             return f"""class {k}:
 {newline.join(enums[k])}"""
+        def add_self_arg(x):
+            return re.sub('\(', '(self, ', x, 1)
         contents = f"""from typing import *
 import enum
 def tree(builder):
@@ -327,7 +330,7 @@ def __ge__(self, other) -> Type: return self
   y = Type()
   z = Type()
   def capture(self, attribute: Type, **kwargs) -> Callable[[], Type]: return transfer_attribute
-  {(newline + '  ').join(map(lambda x: x.replace('(', '(self, '), filter(lambda x: x.startswith('def'), symbols)))}
+  {(newline + '  ').join(map(add_self_arg, filter(lambda x: x.startswith('def'), symbols)))}
   
 {newline.join(map(type_symbol, Type.__subclasses__()))}
 {newline.join(map(enum_namespace, enums.keys()))}

api/static/input_group.py

@@ -1,4 +1,13 @@
-class InputGroup:
+class _InputGroupMeta(type):
+    def __getitem__(cls, args):
+        if isinstance(args, str):
+            class PrefixedInputGroup(InputGroup):
+                prefix = args
+            PrefixedInputGroup.__annotations__ = cls.__annotations__
+            return PrefixedInputGroup
+        return cls
+
+class InputGroup(metaclass=_InputGroupMeta):
     """
     A group of inputs that will be expanded in the node tree.
     

api/tree.py

@@ -47,11 +47,12 @@ def validate_param(param):
                 raise Exception(f"Type of tree input '{param.name}' is not a valid 'Type' subclass.")
         for param in signature.parameters.values():
             if issubclass(param.annotation, InputGroup):
+                prefix = (param.annotation.prefix + "_") if hasattr(param.annotation, "prefix") else ""
                 for group_param, annotation in param.annotation.__annotations__.items():
-                    inputs[group_param] = (annotation, inspect.Parameter.empty, param.name)
+                    inputs[prefix + group_param] = (annotation, inspect.Parameter.empty, param.name, prefix)
             else:
                 validate_param(param)
-                inputs[param.name] = (param.annotation, param.default, None)
+                inputs[param.name] = (param.annotation, param.default, None, None)
 
         # Create the input sockets and collect input values.
         for i, node_input in enumerate(node_group.inputs):
@@ -72,7 +73,7 @@ def validate_param(param):
             if arg[1][2] is not None:
                 if arg[1][2] not in builder_inputs:
                     builder_inputs[arg[1][2]] = signature.parameters[arg[1][2]].annotation()
-                setattr(builder_inputs[arg[1][2]], arg[0], arg[1][0](group_input_node.outputs[i]))
+                setattr(builder_inputs[arg[1][2]], arg[0].replace(arg[1][3], ''), arg[1][0](group_input_node.outputs[i]))
             else:
                 builder_inputs[arg[0]] = arg[1][0](group_input_node.outputs[i])
 
@@ -132,7 +133,10 @@ def group_reference(*args, **kwargs):
             group_outputs = []
             for group_output in result._socket.node.outputs:
                 group_outputs.append(Type(group_output))
-            return tuple(group_outputs)
+            if len(group_outputs) == 1:
+                return group_outputs[0]
+            else:
+                return tuple(group_outputs)
         return group_reference
     if isinstance(name, str):
         return build_tree

api/types.py

@@ -2,6 +2,7 @@
 from bpy.types import NodeSocketStandard
 import nodeitems_utils
 from .state import State
+import geometry_script
 
 def map_case_name(i):
     return ('_' if not i.identifier[0].isalpha() else '') + i.identifier.replace(' ', '_').upper()
@@ -16,7 +17,16 @@ def socket_type_to_data_type(socket_type):
             return socket_type
 
 # The base class all exposed socket types conform to.
-class Type:
+class _TypeMeta(type):
+    def __getitem__(self, args):
+        for s in filter(lambda x: isinstance(x, slice), args):
+            if (isinstance(s.start, float) or isinstance(s.start, int)) and (isinstance(s.stop, float) or isinstance(s.stop, int)):
+                print(f"minmax: ({s.start}, {s.stop})")
+            elif isinstance(s.start, str):
+                print(f"{s.start} = {s.stop}")
+        return self
+
+class Type(metaclass=_TypeMeta):
     socket_type: str
 
     def __init__(self, socket: bpy.types.NodeSocket = None, value = None):
@@ -41,14 +51,10 @@ def __init__(self, socket: bpy.types.NodeSocket = None, value = None):
         self.socket_type = type(socket).__name__
 
     def _math(self, other, operation, reverse=False):
-        math_node = State.current_node_tree.nodes.new('ShaderNodeVectorMath' if self._socket.type == 'VECTOR' else 'ShaderNodeMath')
-        math_node.operation = operation
-        State.current_node_tree.links.new(self._socket, math_node.inputs[1 if reverse else 0])
-        if issubclass(type(other), Type):
-            State.current_node_tree.links.new(other._socket, math_node.inputs[0 if reverse else 1])
+        if self._socket.type == 'VECTOR':
+            return geometry_script.vector_math(operation=operation, vector=(other, self) if reverse else (self, other))
         else:
-            math_node.inputs[0 if reverse else 1].default_value = other
-        return Type(math_node.outputs[0])
+            return geometry_script.math(operation=operation, value=(other, self) if reverse else (self, other))
 
     def __add__(self, other):
         return self._math(other, 'ADD')
@@ -81,30 +87,19 @@ def __rmod__(self, other):
         return self._math(other, 'MODULO', True)
 
     def _compare(self, other, operation):
-        compare_node = State.current_node_tree.nodes.new('FunctionNodeCompare')
-        compare_node.data_type = 'FLOAT' if self._socket.type == 'VALUE' else self._socket.type
-        compare_node.operation = operation
-        a = None
-        b = None
-        for node_input in compare_node.inputs:
-            if not node_input.enabled:
-                continue
-            elif a is None:
-                a = node_input
-            else:
-                b = node_input
-        State.current_node_tree.links.new(self._socket, a)
-        if issubclass(type(other), Type):
-            State.current_node_tree.links.new(other._socket, b)
-        else:
-            b.default_value = other
-        return Type(compare_node.outputs[0])
+        return geometry_script.compare(operation=operation, a=self, b=other)
 
     def __eq__(self, other):
-        return self._compare(other, 'EQUAL')
+        if self._socket.type == 'BOOLEAN':
+            return self._boolean_math(other, 'XNOR')
+        else:
+            return self._compare(other, 'EQUAL')
 
     def __ne__(self, other):
-        return self._compare(other, 'NOT_EQUAL')
+        if self._socket.type == 'BOOLEAN':
+            return self._boolean_math(other, 'XOR')
+        else:
+            return self._compare(other, 'NOT_EQUAL')
 
     def __lt__(self, other):
         return self._compare(other, 'LESS_THAN')
@@ -118,6 +113,44 @@ def __gt__(self, other):
     def __ge__(self, other):
         return self._compare(other, 'GREATER_EQUAL')
 
+    def _boolean_math(self, other, operation, reverse=False):
+        boolean_math_node = State.current_node_tree.nodes.new('FunctionNodeBooleanMath')
+        boolean_math_node.operation = operation
+        a = None
+        b = None
+        for node_input in boolean_math_node.inputs:
+            if not node_input.enabled:
+                continue
+            elif a is None:
+                a = node_input
+            else:
+                b = node_input
+        State.current_node_tree.links.new(self._socket, a)
+        if other is not None:
+            if issubclass(type(other), Type):
+                State.current_node_tree.links.new(other._socket, b)
+            else:
+                b.default_value = other
+        return Type(boolean_math_node.outputs[0])
+    
+    def __and__(self, other):
+        return self._boolean_math(other, 'AND')
+
+    def __rand__(self, other):
+        return self._boolean_math(other, 'AND', reverse=True)
+    
+    def __or__(self, other):
+        return self._boolean_math(other, 'OR')
+    
+    def __ror__(self, other):
+        return self._boolean_math(other, 'OR', reverse=True)
+    
+    def __invert__(self):
+        if self._socket.type == 'BOOLEAN':
+            return self._boolean_math(None, 'NOT')
+        else:
+            return self._math(-1, 'MULTIPLY')
+    
     def _get_xyz_component(self, component):
         if self._socket.type != 'VECTOR':
             raise Exception("`x`, `y`, `z` properties are not available on non-Vector types.")