Custom Nodes

nodes.py

import bpy
import re
from bpy.types import NodeTree, Node, NodeSocket

import nodeitems_utils
from nodeitems_utils import NodeCategory, NodeItem

# Implementation of custom nodes from Python
from bpy import types
from contextlib import contextmanager


class IO:
    types = ()
    sockets = {}
    
    def __init__(self, data_type):
        self._data_type = data_type  
    
    @property
    def socket_name(self):
        return self.__class__.sockets[self._data_type]
           
    @classmethod
    def add_socket_type(cls, name, type_):
        cls.sockets[type_] = name
        cls.types = tuple(cls.sockets.keys())


class Input(IO):
    pass


class Output(IO):
    pass


# Register some node types
IO.add_socket_type('NodeSocketFloat', float)
IO.add_socket_type('NodeSocketInt', int)


class NodeProxyBase(types.Node):
    """Mediator class (handles user-defined data and maps to bpy nodes"""
    
    def init(self, context):
        with forbid_compiling():
            inputs = self.inputs
            for name, io_pin in self._inputs.items():
                inputs.new(io_pin.socket_name, name)
                print(name, io_pin.socket_name)
                
            outputs = self.outputs
            for name, io_pin in self._outputs.items():
                outputs.new(io_pin.socket_name, name)
                

def camel_case_to_name(value):
    s1 = re.sub('(.)([A-Z][a-z]+)', r'\1 \2', value)
    return re.sub('([a-z0-9])([A-Z])', r'\1 \2', s1)


def node_builder(cls):
    """Builds bpy node from cls"""
    attrs = cls.__dict__.copy()
    name = cls.__name__
    bases = cls.__bases__
    
    io_pins = {k: v for k, v in attrs.items() if isinstance(v, IO)}
    inputs = {k: v for k, v in io_pins.items() if isinstance(v, Input)}
    outputs = {k: v for k, v in io_pins.items() if not k in inputs}
    
    attrs['_inputs'] = inputs
    attrs['_outputs'] = outputs
    
    attrs['bl_label'] = attrs.get("label", camel_case_to_name(name))
    attrs['bl_icon'] = attrs.get('icon', 'NONE')

    # Remove pins from attr dict
    for k in io_pins:
        attrs.pop(k)
    
    bases = bases[1:] + (NodeProxyBase,)
    cls = type(name, bases, attrs)
    return cls


@contextmanager
def forbid_compiling():
    # Stop compiling here
    yield
    # Resume compiler here
    

class AnimationNode:
    
    @classmethod
    def find_subclasses(cls, instigator=None):
        if instigator is None:
           instigator = cls
                
        for subcls in cls.__subclasses__():
            yield from subcls.find_subclasses(instigator)
        
        if instigator is cls:
           return
        
        yield cls


class MyCustomTree(NodeTree):
    bl_label = 'Custom Node Tree'
    bl_icon = 'NODETREE'
    
    
# Custom nodes finally declared!
class OtherNode(AnimationNode):
    '''A custom node'''
    
    x = Input(int)
    y = Input(int)
    
    z = Output(int)


class OtherNode2(AnimationNode):
    '''A custom node'''
    
    x = Input(int)
    y = Input(int)
    z = Input(float)
    
    p = Output(float)
    

built_nodes = [node_builder(n) for n in AnimationNode.find_subclasses()]
node_categories = [NodeCategory("SOMENODES", "Some Nodes", items=[NodeItem(cls.__name__) for cls in built_nodes])]

def register():
    bpy.utils.register_class(MyCustomTree)
    for cls in built_nodes:
        bpy.utils.register_class(cls)

    nodeitems_utils.register_node_categories("CUSTOM_NODES", node_categories)


def unregister():
    nodeitems_utils.unregister_node_categories("CUSTOM_NODES")

    bpy.utils.unregister_class(MyCustomTree)
    for cls in built_nodes:
        bpy.utils.unregister_class(cls)


if __name__ == "__main__":
    try:
        unregister()
    except:
        pass
     
    register()