node_kit/core/_node.py

import bpy
from mathutils import Color, Vector

from .sockets import Input, Output


class Node:
    """Blender Node abstraction."""

    def __init__(self, bl_node, parent):

        self.bl_node = bl_node
        self.tree = parent
        self.id = hex(id(self.bl_node))

        self.data = {}
        self.parameters = []

        self._parent = None
        self._scene = None

        for prop in self.bl_node.bl_rna.properties:
            if prop.is_readonly:
                continue

            prop_id = prop.identifier

            setattr(self, prop_id, getattr(self.bl_node, prop_id))
            self.parameters.append(prop_id)

        self.inputs = [Input(ipt, self.tree) for ipt in self.bl_node.inputs]
        self.outputs = [Output(opt, self.tree) for opt in self.bl_node.outputs]

    @property
    def parent(self):
        """Get the Node from all the other nodes in the tree checking that the
        parent of its blender node is the same as the blender node we are comparing.

        Returns:
            Node: Node parent.
        """

        if self._parent:
            return self._parent

        # if blender node doesn't have a parent
        if not self.bl_node.parent:
            self._parent = None
            return self._parent

        for node in self.tree.nodes:
            if node.bl_node == self.bl_node.parent:
                self._parent = node
                return self._parent

    @parent.setter
    def parent(self, value):
        """Set the Node parent, using the python object, it's id or the blender node.

        Args:
            value (Node|str|bpy.types.Node): Node, id or blender node to set as parent.
        """

        # Node object case
        if isinstance(value, Node):
            self._parent = value

        # Node id case
        elif isinstance(value, str) and value.startswith('0x'):
            for node in self.tree.nodes:
                if node.id == value:
                    self._parent = node
            else:
                print('Cannot find parent')

        # blender node case
        elif isinstance(value, bpy.types.Node):
            for node in self.tree.nodes:
                if node.bl_node == value:
                    self._parent = node

        if self._parent:
            self.bl_node.parent = self._parent.bl_node

    @classmethod
    def from_blender_node(cls, bl_node, tree):
        """Instanciate an abstract class based of the blender node idname.

        Args:
            bl_node (bpy.types.Node): Blender Node To create abstraction from.
            tree (NodeTree): Node tree object node belongs to.

        Returns:
            Node: Node abstract according to the blender node type.
        """
        if bl_node.bl_idname == 'CompositorNodeRLayers':
            return RenderLayersNode(bl_node, tree)

        elif bl_node.bl_idname == 'CompositorNodeValToRGB':
            return ColorRampNode(bl_node, tree)

        else:
            return cls(bl_node, tree)

    @classmethod
    def from_dict(cls, data, tree):
        """Create all nodes from their dict representation.

        Args:
            data (dict): dict nodes representation.
            tree (Tree): blender node tree abstraction.

        Returns:
            Node: Create abstract node.
        """

        new_bl_node = tree.bl_node_tree.nodes.new(type=data['bl_idname'])
        node = cls.from_blender_node(new_bl_node, tree)

        node.id = data['id']
        for p in node.parameters:
            setattr(node, p, data[p])

            # set attribute on the blender node only if correct type is retrieve
            if p not in ('parent', 'scene'):
                setattr(node.bl_node, p, getattr(node, p))

        node.inputs = [Input.from_dict(ipt_data, node) for ipt_data in data['inputs'].values()]
        node.outputs = [Output.from_dict(opt_data, node) for opt_data in data['outputs'].values()]
        return node

    def dump(self):
        """Export currrent Node to its dict representation.

        Returns:
            dict: Node dict representation.
        """

        for prop_id in self.parameters:

            if not hasattr(self, prop_id):
                continue

            attr_value = getattr(self, prop_id)
            if attr_value is None:
                attr_value = None

            elif isinstance(attr_value, Node):
                attr_value = attr_value.id

            elif isinstance(attr_value, (Color, Vector)):
                attr_value = list(attr_value)

            self.data[prop_id] = attr_value

        self.data['id'] = self.id
        self.data['inputs'] = {ipt.id: ipt.dump() for ipt in self.inputs}
        self.data['outputs'] = {opt.id: opt.dump() for opt in self.outputs}

        return self.data


class RenderLayersNode(Node):
    """Blender Render Layers Node abstraction"""

    @property
    def scene(self):
        """Get the name of the scene used by the node.

        Returns:
            str: scene name.
        """
        if self._scene:
            return self._scene.name

    @scene.setter
    def scene(self, value):
        """Set the blender scene using the bpy Scene object or its name.

        Args:
            value (str|bpy.types.Scene): scene name or scene object to set the node.
        """
        if isinstance(value, str):
            self._scene = bpy.data.scenes[value]

        elif isinstance(value, bpy.types.Scene):
            self._scene = value

        if self._scene:
            self.bl_node.scene = self._scene


class Link:
    """Blender Link abstraction."""

    def __init__(self, bl_link, parent):

        self.bl_link = bl_link
        self.tree = parent
        self.id = hex(id(self.bl_link))

        self.input = self.bl_link.to_socket
        self.output = self.bl_link.from_socket

        self.data = {}

    def dump(self):

        self.data['id'] = self.id

        return self.data