node_kit/core/node_tree.py

import json
from pathlib import Path

from .node import Node, Link


class NodeTree:
    """Blender node tree abstraction."""

    def __init__(self, bl_node_tree):

        self.bl_node_tree = bl_node_tree

        self.data = {}

        self.links = [Link(lnk, parent=self) for lnk in self.bl_node_tree.links]
        self.nodes = []
        for n in self.bl_node_tree.nodes:
            self.nodes.append(Node.from_blender_node(n, self))

    def to_dict(self, select_only=False):
        """Convert all blender nodes and links inside the tree into a dictionnary.

        Args:
            select_only (bool, optional): True to convert only selected nodes.
                Defaults to False.

        Returns:
            dict: Nodes and links as dict. 
        """
        self.data['nodes'] = {n.id: n.to_dict() for n in self.nodes if not select_only or (select_only and n.select)}
        self.data['links'] = [l.id for l in self.links]

        return self.data

    def ingest_dict(self, data):
        """From a Tree dict representation, create new nodes with their attributes.
        Then create a connection dict by comparing link id from inputs and outputs of each nodes.
        Use this dict to link nodes between each others.

        Args:
            data (dict): Tree dict representation to generate nodes and links from.
        """

        connections = {}

        self.data = data

        for node_id, node_data in self.data['nodes'].items():

            new_node = Node.from_dict(node_data, self)
            self.nodes.append(new_node)

            new_node.bl_node.select = True

            for ipt in new_node.inputs:
                if ipt.is_linked:
                    connections.setdefault(ipt.link, {})['input'] = ipt.bl_input

            for opt in new_node.outputs:
                if opt.is_linked:
                    for link in opt.link:
                        connections.setdefault(link, {})['output'] = opt.bl_output

        for link_id in self.data['links']:
            ipt = connections[link_id]['input']
            opt = connections[link_id]['output']

            self.bl_node_tree.links.new(ipt, opt)