node_kit/core/dumper.py

import bpy
import mathutils
from pprint import pprint
import json
import itertools
from copy import copy
from os.path import abspath



def get_default(prop):
    """Get the default value of a bl property"""

    if getattr(prop, 'is_array', False):
        return list(prop.default_array)
    elif hasattr(prop, 'default'):
        return prop.default
        

def get_dumper(bl_object, fallback=None):
    """Find the right dumper type by checking inheritance"""
    for dp in dumpers:                
        if isinstance(bl_object, dp.bl_type):
            return dp
    
    return fallback or Dumper


def get_bl_object(data):
    """Find the bl object for loading data into it depending on the type and the context"""
    if data.get('_new', {}).get('type') == 'GeometryNodeTree':
        return bpy.context.object.modifiers.active.node_group


def dump(ob):
    """Generic Recursive Dump, convert any object into a dict"""
    Dumper.pointers.clear()
    
    if isinstance(ob, (list, tuple)):
        data = [get_dumper(o).dump(o) for o in ob]
    else:
        data = get_dumper(ob).dump(ob)
    
    Dumper.pointers.clear()

    return data
    

def load(data, bl_object=None):
    """Generic Load to create an object from a dict"""

    Dumper.pointers.clear()
    #print(Dumper.pointers)

    if bl_object is None:
        bl_object = get_bl_object(data)
    
    dumper = get_dumper(bl_object)
    dumper.load(data, bl_object)

    Dumper.pointers.clear()


def set_attribute(bl_object, attr, value):
    try:
        setattr(bl_object, attr, value)
    except Exception as e:
        print(e)


class Dumper:
    pointers = {}
    includes = []
    excludes = ["rna_type", "bl_rna", 'id_data', 'depsgraph']
    
    @classmethod 
    def properties(cls, bl_object):
        if cls.includes and not cls.excludes:
            return [bl_object.bl_rna.properties[p] for p in cls.includes]
        else:
            return [ p for p in bl_object.bl_rna.properties if not
                     p.identifier.startswith('bl_') and p.identifier not in cls.excludes]
    
    @classmethod
    def new(cls, data):
        print(f'New not implemented for data {data}')

    @classmethod  
    def load(cls, data, bl_object=None):
        if bl_object is None:
            bl_object = cls.new(data)

        if bl_object is None:
            return
        
        #pprint(data)
        if bl_pointer := data.get('bl_pointer'):
            cls.pointers[bl_pointer] = bl_object

        props = cls.properties(bl_object)
        for key, value in sorted(data.items(), key=lambda x: props.index(x[0]) if x[0] in props else 0):
            if key.startswith('_') or key not in bl_object.bl_rna.properties:
                continue
            
            prop = bl_object.bl_rna.properties[key]
            attr = getattr(bl_object, key)
                                    
            if prop.type == 'COLLECTION':
                dumper = PropCollection
                if hasattr(attr, 'bl_rna'): 
                    bl_type = attr.bl_rna.type_recast()
                    dumper = get_dumper(bl_type, fallback=PropCollection)
                
                dumper.load(value, attr)
                continue
                
            elif prop.type == 'POINTER':
                # if key == 'node_tree':
                #     print('--------------')
                #     print(bl_object, value)
                #     print(cls.pointers)

                if isinstance(value, int): # It's a pointer
                    value = cls.pointers[value]
                elif value is None:
                    set_attribute(bl_object, key, value)

                else:
                    bl_type = prop.fixed_type.bl_rna.type_recast()
                    dumper = get_dumper(bl_type)

                    # If the pointer exist register the pointer then load data
                    #print('-----', value)
                    #pointer = 
                    if attr is None:
                        attr = dumper.new(value)

                    dumper.load(value, attr)
                    #attr = getattr(bl_object, key)
                    #if not attr:
                    cls.pointers[value['bl_pointer']] = attr

                    if hasattr(attr, 'update'):
                        attr.update()
                
                    value = attr

                if not prop.is_readonly:
                    set_attribute(bl_object, key, value)
                
                # Some coll needs a manual update like curve mapping
                if hasattr(attr, 'update'):
                    attr.update()
            
            elif not prop.is_readonly:
                #print(key, value)
                set_attribute(bl_object, key, value)
                continue
        
        #return bl_object

    @classmethod
    def dump(cls, bl_object):
        if isinstance(bl_object, (str, int, float, dict, list, type(None))):
            return bl_object
        
        #print('Dumping object', bl_object)

        data = {"bl_pointer": bl_object.as_pointer()}
        cls.pointers[bl_object.as_pointer()] = bl_object


        for prop in cls.properties(bl_object):            
            if not hasattr(bl_object, prop.identifier):
                print(f'{bl_object} has no attribute {prop.identifier}')
                continue
            
            #print(prop.identifier)

            value = getattr(bl_object, prop.identifier)
            
            # Not storing default value
            if prop.identifier not in cls.includes:
                if (array := getattr(prop, 'default_array', None)) and value == array:
                    continue
                if isinstance(value, (str, int, float)) and value == prop.default:
                    continue
                   
            if getattr(prop, "is_array", False):
                value = PropArray.dump(value)
            
            elif prop.type == 'COLLECTION':
                value = PropCollection.dump(value)
                
            elif prop.type == 'POINTER' and value:
                #if prop.identifier == 'image':
                #    print(bl_object, cls.pointers)
                if value.as_pointer() in cls.pointers:
                    value = value.as_pointer()
                else:
                    # print('Register Pointer', value.as_pointer(), value)
                    cls.pointers[value.as_pointer()] = value
                    # print(cls.pointers)
                    # print()
                    dumper = get_dumper(value)
                    value = dumper.dump(value)
                                    
            elif bl_object.is_property_readonly(prop.identifier):
                continue
            
            else:
                dumper = get_dumper(value)
                value = dumper.dump(value)
            
            data[prop.identifier] = value
        
        return data
        

class PropCollection(Dumper):
    bl_type = bpy.types.bpy_prop_collection
    
    @classmethod
    def dump(cls, coll):
        if not len(coll):
            return []
        
        dumper = get_dumper(coll[0])
        values = [dumper.dump(e) for e in coll]

        # Value cannot be None
        return [v for v in values if v is not None]
        
    
    @classmethod
    def load(cls, values, coll):
        if not values:
            return

        dumper = None
        
        if not hasattr(coll, 'new'): # Static collection
            for item, value in zip(coll, values):
                dumper = dumper or get_dumper(item)
                dumper.load(value, item)

            return     

        new_func = coll.bl_rna.functions['new']
        for i, value in enumerate(values):
                            
            if value.get('_new'):
                params = value['_new']
            else:
                params = {k: value.get(k, get_default(v)) for k, v in new_func.parameters.items()[:-1]}
            
            # Replace arg pointer with bl object
            valid_pointers = True
            for param in coll.bl_rna.functions['new'].parameters:
                if param.identifier not in params or param.type != 'POINTER':
                    continue
                
                pointer_id = params[param.identifier]
                if bl_object := cls.pointers.get(pointer_id):
                    params[param.identifier] = bl_object
                else:
                    print(f'No Pointer found for param {param.identifier} of {coll}')
                    valid_pointers = False
            
            if not valid_pointers:
                continue

                #print(param.identifier, cls.pointers[pointer_id])

            try:

                item = coll.new(**params)
            except RuntimeError as e:
                #print(e, coll.data)
                #print()
                try:
                    item = coll[i]
                except IndexError as e:
                    #print(e, coll.data)
                    break
            
            dumper = get_dumper(item)    
            dumper.load(value, item)#(item, value)
    
    
class PropArray(Dumper):
    bl_type = bpy.types.bpy_prop_array
        
    @classmethod
    def dump(cls, array):
        flat_array = []
        for item in array:
            if isinstance(item, (int, float)):
                flat_array.append(item)
            else:
                flat_array.extend(cls.dump(item))
        return flat_array


class NodeSocket(Dumper):
    bl_type = bpy.types.NodeSocket
    excludes = Dumper.excludes + ["node", "links", "display_shape", "rna_type", "link_limit"]

    @classmethod  
    def dump(cls, socket):
        if socket.is_unavailable:
            return None

        #cls.pointers[socket.as_pointer()] = socket
        
        data = super().dump(socket)
                
        #data["_id"] = socket.as_pointer()
        #data.pop('name', '')
        
        return data


class NodeGeometryRepeatOutputItems(PropCollection):
    bl_type = bpy.types.NodeGeometryRepeatOutputItems

    @classmethod
    def load(cls, values, coll):
        coll.clear()

        super().load(values, coll)


class NodeLink(Dumper):
    bl_type = bpy.types.NodeLink
    
    @classmethod  
    def dump(cls, link):
        return {"_new": {
                    "input": link.from_socket.as_pointer(), 
                    "output": link.to_socket.as_pointer()
                    }
                }
        

class NodeTreeInterfaceSocket(Dumper):
    bl_type = bpy.types.NodeTreeInterfaceSocket
    excludes = Dumper.excludes + ["parent", "interface_items"]
    
    @classmethod  
    def dump(cls, socket):
        #cls.pointers[socket.as_pointer()] = socket
        
        data = super().dump(socket)
        #data["_id"] = socket.as_pointer()
        
        data['_new'] = {"name": data.get('name', '')}

        if socket.item_type == 'SOCKET':
            data['_new']["in_out"] = socket.in_out

        
        # It's a real panel not the interface root
        if socket.parent.parent: 
            data['parent'] = socket.parent.as_pointer()
        
        return data
    

class NodeSockets(PropCollection):

    @classmethod  
    def load(cls, values, coll):

        #return
        
        node_sockets = [s for s in coll if not s.is_unavailable]
        for socket, value in zip(node_sockets, values):
            cls.pointers[value['bl_pointer']] = socket
            Dumper.load(value, socket)
            # for k, v in value.items():
            #     if k not in socket.bl_rna.properties:
            #         continue
            #     setattr(socket, k, v)

        """
        # Match Inputs Pointers
        node_sockets = [s for s in coll if not s.is_unavailable]
        if len(node_sockets) == len(inputs): # Match by index
            super().load({"inputs": inputs}, node)
            for socket, value in zip(node_sockets, coll):
                cls.pointers[value['_id']] = socket
        else: # Match by name
            print(f'Match Inputs by Name for node {node}')
            for socket in node_sockets:
                index = next((i for i, v in enumerate(inputs) if v['name'] == socket.name), None)
                if index is None:
                    continue

                value = inputs[index]
                print(socket, value)
                cls.pointers[value['_id']] = socket

                Dumper.load(value, socket)
                del inputs[index]
        """


class NodeInputs(NodeSockets):
    bl_type = bpy.types.NodeInputs


class NodeOutputs(NodeSockets):
    bl_type = bpy.types.NodeOutputs


class Node(Dumper):
    bl_type = bpy.types.Node
    excludes = Dumper.excludes + ["dimensions", "height", "internal_links", "paired_output"]

    @classmethod  
    def dump(cls, node=None):
        #cls.pointers[node.as_pointer()] = node
        
        data = super().dump(node)
        #data["_id"] = node.as_pointer()
        data["_new"] = {"type": node.bl_rna.identifier} # 'node_tree': node.id_data.as_pointer()

        if paired_output := getattr(node, "paired_output", None):
            data["_pair_with_output"] = paired_output.as_pointer()
        
        #if node.parent:
        #    data['location'] -= Vector()node.parent.location

        return data

    @classmethod  
    def load(cls, data, node):
        if node is None:
            return
        #cls.pointers[data['bl_pointer']] = node
        
        inputs = copy(data.pop('inputs', []))
        outputs = copy(data.pop('outputs', []))

        super().load(data, node)

        data['inputs'] = inputs
        data['outputs'] = outputs

        # Loading input and outputs after the properties
        super().load({"inputs": inputs, "outputs": outputs}, node)

        if node.parent:
            node.location += node.parent.location
        
        #if node.type != 'FRAME':
        #    node.location.y -= 500
        

class CompositorNodeGlare(Node):
    bl_type = bpy.types.CompositorNodeGlare

    includes = ["quality"]


class NodeTreeInterface(Dumper):
    bl_type = bpy.types.NodeTreeInterface
    
    @classmethod  
    def load(cls, data, interface):    

        print('Load Interface')

        for value in data.get('items_tree', []):
            item_type = value.get('item_type', 'SOCKET')
            if item_type == 'SOCKET':
                item = interface.new_socket(**value['_new'])
            elif item_type == 'PANEL':
                #print(value['_new'])
                item = interface.new_panel(**value['_new'])
            
            NodeTreeInterfaceSocket.load(value, item)
    
        interface.active_index = data.get('active_index', 0)


class Nodes(PropCollection):
    bl_type = bpy.types.Nodes

    @classmethod
    def load(cls, values, coll):
        super().load(values, coll)

        # Pair zone input and output
        for node_data in values:
            if paired_output_id := node_data.get('_pair_with_output', None):
                node = cls.pointers[node_data['bl_pointer']]
                node.pair_with_output(cls.pointers[paired_output_id])

                #print(node, node_data['outputs'])

                Dumper.load({"inputs": node_data['inputs'], "outputs": node_data['outputs']}, node)


class NodeTree(Dumper):
    bl_type = bpy.types.NodeTree
    excludes = []
    includes = ["name", "interface", "nodes", "links"]

    @classmethod  
    def new(cls, data):
        if link := data.get('_link'):
            with bpy.data.libraries.load(link['filepath'], link=True) as (data_from, data_to):
                setattr(data_to, link['data_type'], [link['name']])
            return getattr(data_to, link['data_type'])[0]

        return bpy.data.node_groups.new(**data["_new"])

    @classmethod  
    def dump(cls, node_tree):
        if node_tree.library:
            data = {'bl_pointer': node_tree.as_pointer()}
            filepath = abspath(bpy.path.abspath(node_tree.library.filepath, library=node_tree.library.library))
            data["_link"] = {"filepath": filepath, "data_type": 'node_groups', 'name': node_tree.name}
        else:
            data = super().dump(node_tree)
            data["_new"] = {"type": node_tree.bl_rna.identifier, 'name': node_tree.name}
        
        return data
    
    
class Points(PropCollection):

    @classmethod
    def load(cls, values, coll):
        new_func = coll.bl_rna.functions['new']
        params = {k: get_default(v)+1.1 for k, v in new_func.parameters.items()[:-1]}
                
        # Match the same number of elements in collection
        if len(values) > len(coll):
            for _ in range(len(values) - len(coll)):
                coll.new(**params)
        
        for i, value in enumerate(values):
            Dumper.load(value, coll[i])
            #for k, v in value.items():
                #setattr(coll[i], k, v)
                

class CurveMapPoints(Points):
    bl_type = bpy.types.CurveMapPoints
    

class ColorRampElements(Points):
    bl_type = bpy.types.ColorRampElements


class CompositorNodeOutputFileLayerSlots(PropCollection):
    bl_type = bpy.types.CompositorNodeOutputFileLayerSlots

    @classmethod
    def load(cls, values, coll):
        coll.clear()

        super().load(values, coll)


class CompositorNodeOutputFileFileSlots(PropCollection):
    @classmethod
    def load(cls, values, coll):
        coll.clear()

        super().load(values, coll)


class AOVs(PropCollection):
    bl_type = bpy.types.AOVs

    @classmethod
    def load(cls, values, coll):
        for value in values:
            aov = coll.get(value['name'])

            if not aov:
                aov = coll.add()
            
            Dumper.load(value, aov)


class Image(Dumper):
    bl_type = bpy.types.Image

    excludes = []
    includes = ['name', 'filepath']

    @classmethod
    def new(cls, data):
        # image = next((  img for img in bpy.data.images if not img.library 
        #                 and img.filepath == data['filepath']), None)

        # if image is None:
        #     image = bpy.data.images.load(data['filepath'])

        return bpy.data.images.load(data['filepath'], check_existing=True)


class Material(Dumper):
    bl_type = bpy.types.Material

    excludes = Dumper.excludes + ['preview', "original"]

    @classmethod
    def new(cls, data):
        material = bpy.data.materials.get(data.get('name', ''))

        if material is None:
            material = bpy.data.materials.new(data['name'])
        
        return material


class Object(Dumper):
    bl_type = bpy.types.Object
    excludes = []
    includes = ['name']

    @classmethod
    def new(cls, data):
        if name := data.get('name'):
            return bpy.data.objects.get(name)


class Scene(Dumper):
    bl_type = bpy.types.Scene
    excludes = []
    includes = ['name']
    

    @classmethod
    def new(cls, data):
        if scene := bpy.data.scenes.get(data.get('name', '')):
            return scene
        
        return bpy.data.scenes.new(name=data.get('name', ''))

    """
    @classmethod  
    def dump(cls, scene):
        view_layer = scene.view_layers[node.layer]
        view_layer_data = ViewLayer.dump(view_layer)

        return {
            'bl_pointer': scene.as_pointer(),
            'name': scene.name,
            'render' : {'bl_pointer': scene.render.as_pointer(), "engine": scene.render.engine},
            'view_layers': [view_layer_data]
        }
        """

class Collection(Dumper):
    bl_type = bpy.types.Collection
    includes = ['name']
    excludes = []

    @classmethod
    def new(cls, data):
        if name := data.get('name'):
            return bpy.data.collections.get(name)

    # @classmethod
    # def dump(cls, data):
    #     data = super().dump(data)

    #     data['render'] = {"engine": scene.render.engine}
    #     return data


class CompositorNodeRLayers(Node):
    bl_type = bpy.types.CompositorNodeRLayers

    excludes = Dumper.excludes + ['scene']

    @classmethod  
    def load(cls, data, node):

        #print('load CompositorNodeRLayers')

        scene_data = data.pop('scene')
        #print(scene_data)
        layer = data.pop('layer')
        scene = Scene.new(scene_data)
        Scene.load(scene_data, scene)
        
        node.scene = scene
        node.layer = layer

        super().load(data, node)

        # Resetter the view_layer because it might have been created
        # with the scene attr in the dictionnary and nor available yet

        #print(bpy.)



    @classmethod  
    def dump(cls, node):
        # Add scene and viewlayer passes
        data = super().dump(node)

        #if 

        view_layer = node.scene.view_layers[node.layer]
        view_layer_data = ViewLayer.dump(view_layer)

        '''
        view_layer_data = {
            "name": view_layer.name}
        properties = {p.name: p for p in view_layer.bl_rna.properties}
        for prop in view_layer.bl_rna:
            if prop.identifier.startswith('use_pass'):
                view_layer_data[prop.identifier]
                '''

        #cls.pointers[bl_object.as_pointer()] = bl_object

        data['scene'] = {
            'bl_pointer': node.scene.as_pointer(),
            'name': node.scene.name,
            'render' : {'bl_pointer': node.scene.render.as_pointer(), "engine": node.scene.render.engine},
            'view_layers': [view_layer_data]
        }

        return data


class ViewLayer(Dumper):
    bl_type = bpy.types.ViewLayer
    excludes = Dumper.excludes + ['freestyle_settings', 'eevee', 'cycles', 'active_layer_collection',
        'active_aov', 'active_lightgroup_index', 'layer_collection', 'lightgroups', 'material_override',
        'objects', 'use']
    #includes = ['name']
    

class ViewLayers(PropCollection):
    bl_type = bpy.types.ViewLayers

    @classmethod
    def load(cls, values, coll):
        #print('LOAD VIEWLAYERS', values)
        for value in values:
            view_layer = coll.get(value['name'])

            if view_layer is None:
                view_layer = coll.new(value['name'])
            
            Dumper.load(value, view_layer)


dumpers = [
    CompositorNodeRLayers,
    CompositorNodeGlare,
    Node,
    NodeSocket,
    NodeTree,
    NodeLink,
    NodeTreeInterface,
    NodeTreeInterfaceSocket,
    NodeGeometryRepeatOutputItems,
    Image,
    Material,
    Object,
    Scene,
    Collection,
    ViewLayer,
    CurveMapPoints,
    ColorRampElements,
    NodeInputs,
    NodeOutputs,
    Nodes,
    ViewLayers,
    PropCollection,
    AOVs,
    PropArray,
    CompositorNodeOutputFileLayerSlots,
    CompositorNodeOutputFileFileSlots,
]