Improve code generation and documentation

types.py

@@ -0,0 +1,115 @@
+import bpy
+from bpy.types import NodeSocketStandard
+import nodeitems_utils
+from .state import State
+
+# The base class all exposed socket types conform to.
+class Type:
+    socket_type: str
+
+    def __init__(self, socket: bpy.types.NodeSocket = None, value = None):
+        if value is not None:
+            input_nodes = {
+                int: ('FunctionNodeInputInt', 'integer'),
+                bool: ('FunctionNodeInputBool', 'boolean'),
+                str: ('FunctionNodeInputString', 'string'),
+                tuple: ('FunctionNodeInputVector', 'vector'),
+                float: ('ShaderNodeValue', None),
+            }
+            if not type(value) in input_nodes:
+                raise Exception(f"'{value}' cannot be expressed as a node.")
+            input_node_info = input_nodes[type(value)]
+            value_node = State.current_node_tree.nodes.new(input_node_info[0])
+            if input_node_info[1] is None:
+                value_node.outputs[0].default_value = value
+            else:
+                setattr(value_node, input_node_info[1], value)
+            socket = value_node.outputs[0]
+        self._socket = socket
+        self.socket_type = type(socket).__name__
+    
+    def _math(self, other, operation):
+        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[0])
+        if issubclass(type(other), Type):
+            State.current_node_tree.links.new(other._socket, math_node.inputs[1])
+        else:
+            math_node.inputs[1].default_value = other
+        return Type(math_node.outputs[0])
+
+    def __add__(self, other):
+        return self._math(other, 'ADD')
+    
+    def __sub__(self, other):
+        return self._math(other, 'SUBTRACT')
+    
+    def __mul__(self, other):
+        return self._math(other, 'MULTIPLY')
+    
+    def __truediv__(self, other):
+        return self._math(other, 'DIVIDE')
+    
+    def __mod__(self, other):
+        return self._math(other, 'MODULO')
+    
+    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
+        State.current_node_tree.links.new(self._socket, compare_node.inputs[0])
+        if issubclass(type(other), Type):
+            State.current_node_tree.links.new(other._socket, compare_node.inputs[1])
+        else:
+            compare_node.inputs[1].default_value = other
+        return Type(compare_node.outputs[0])
+    
+    def __eq__(self, other):
+        return self._compare(other, 'EQUAL')
+    
+    def __ne__(self, other):
+        return self._compare(other, 'NOT_EQUAL')
+    
+    def __lt__(self, other):
+        return self._compare(other, 'LESS_THAN')
+    
+    def __le__(self, other):
+        return self._compare(other, 'LESS_EQUAL')
+    
+    def __gt__(self, other):
+        return self._compare(other, 'GREATER_THAN')
+    
+    def __ge__(self, other):
+        return self._compare(other, 'GREATER_EQUAL')
+
+for standard_socket in list(filter(lambda x: 'NodeSocket' in x, dir(bpy.types))):
+    name = standard_socket.replace('NodeSocket', '')
+    if len(name) < 1:
+        continue
+    globals()[name] = type(name, (Type,), { 'socket_type': standard_socket, '__module__': Type.__module__ })
+    if name == 'Vector':
+        def get_component(component):
+            @property
+            def get(self):
+                separate_node = State.current_node_tree.nodes.new('ShaderNodeSeparateXYZ')
+                State.current_node_tree.links.new(self._socket, separate_node.inputs[0])
+                return Type(separate_node.outputs[component])
+            return get
+        globals()[name].x = get_component(0)
+        globals()[name].y = get_component(1)
+        globals()[name].z = get_component(2)
+    if name == 'Int':
+        class IntIterator:
+            def __init__(self, integer):
+                self.integer = integer
+                self.points = State.current_node_tree.nodes.new('GeometryNodePoints')
+                State.current_node_tree.links.new(self.integer._socket, self.points.inputs[0])
+                self.index = State.current_node_tree.nodes.new('GeometryNodeInputIndex')
+                self._did_iterate = False
+            def __next__(self):
+                if not self._did_iterate:
+                    self._did_iterate = True
+                    return Type(self.index.outputs[0]), Type(self.points.outputs[0])
+                else:
+                    raise StopIteration()
+        globals()[name].__iter__ = lambda self: IntIterator(self)