mantis/xForm_nodes_ui.py

import bpy
from .base_definitions import xFormNode
from bpy.types import Node
from .utilities.utilities import (prRed, prGreen, prPurple, prWhite,
                              prOrange,
                              wrapRed, wrapGreen, wrapPurple, wrapWhite,
                              wrapOrange,)
from .base_definitions import get_signature_from_edited_tree
from .xForm_socket_templates import *

def TellClasses():
    return [
        # xFormNullNode,
        xFormBoneNode,
        xFormArmatureNode,
        xFormGeometryObjectNode,
        xFormObjectInstance,
        xFormCurvePin,
        xFormGetBone,
        ]

def default_traverse(self, socket):
    if (socket == self.outputs["xForm Out"]):
        return self.inputs["Relationship"]
    if (socket == self.inputs["Relationship"]):
        return self.outputs["xForm Out"]
    return None

def check_if_connected(start, end, line):
    started=False
    for path_mantis_node in line:
        prWhite("    ", path_mantis_node.signature)
        if path_mantis_node.signature == start.signature:
            started = True
        elif path_mantis_node.signature == end.signature:
            break
        if started:
            if path_mantis_node.inputs.get("Connected"):
                if path_mantis_node.evaluate_input("Connected") == False:
                    return False
    else:
        return False
    return True


def main_draw_label(self): # this will prefer a user-set label, or return the evaluated name
    if self.label:
        return self.label
    if self.inputs['Name'].display_text:
        return self.inputs['Name'].display_text
    return self.name


def custom_prop_display_update(self, mantis_node = None):
    custom_props = {}
    type_map = {
        'BOOL'   : 'ParameterBoolSocket',
        'INT'    : 'ParameterIntSocket',
        'FLOAT'  : 'ParameterFloatSocket',
        'VECTOR' : 'ParameterVectorSocket',
        'STRING' : 'ParameterStringSocket',
    }
    cant_read=0
    if not mantis_node:
        links = [link for link in self.inputs['Custom Properties'].links]
        links.sort(key=lambda a : -a.multi_input_sort_id)
        for link in links:
            if link.from_node.bl_idname == 'UtilityCustomProperty':
                if (link.from_node.inputs['Name'].is_linked or 
                            link.from_node.inputs['Type'].is_linked):
                    cant_read+=1 # we can't eval this without the Mantis data
                    continue # but we need to count it
                else:
                    custom_props[link.from_node.inputs['Name'].default_value] = \
                        type_map[link.from_node.inputs['Type'].default_value]
            else:
                cant_read+=1
        for name, type in custom_props.items():
            if self.outputs.get(name): continue # already there
            else:
                self.outputs.new(type, name)
    else:
        from .misc_nodes import UtilityCustomProperty
        from .mantis_dataclasses import custom_prop_template
        mantis_node.inputs['Custom Properties'].flush_links() # clean/sort the links
        
        if len(mantis_node.inputs['Custom Properties'].links) != len(self.inputs['Custom Properties'].links):
            return # there is a problem somewhere
        for i, link in enumerate(mantis_node.inputs['Custom Properties'].links):
            from .utilities.node_common import trace_single_line
            node_line, _last_socket = trace_single_line(mantis_node, 'Custom Properties', i)
            if not node_line[-1].prepared:
                return
            custom_prop_data = mantis_node.evaluate_input('Custom Property', i)
            if isinstance(node_line[-1], UtilityCustomProperty) and \
                    not isinstance(custom_prop_data, custom_prop_template):
                return
            elif not isinstance(custom_prop_data, custom_prop_template):
                continue # wrong connection
            name = custom_prop_data.name
            type = custom_prop_data.prop_type
            custom_props[name] = type_map[type]
    from .utilities.utilities import get_socket_maps, do_relink
    _socket_map_in, socket_map_out = get_socket_maps(self)
    for name, type in custom_props.items():
        if socket_map_out.get(name) is None:
            socket_map_out[name] = None # initialize the new ones
    # remove and replace sockets
    if cant_read == 0:
        remove_me = self.outputs[1:]
        while remove_me:
            remove_socket = remove_me.pop()
            self.outputs.remove(remove_socket)
        for name, type in custom_props.items():
            socket = self.outputs.new(type, name)
            do_relink(self, socket, socket_map_out, in_out='OUTPUT')


class xFormBoneNode(Node, xFormNode):
    '''A node representing a Bone'''
    bl_idname = 'xFormBoneNode'
    bl_label = "Bone"
    bl_icon = 'BONE_DATA'
    
    display_ik_settings : bpy.props.BoolProperty(default=False)
    display_vp_settings : bpy.props.BoolProperty(default=False)
    display_def_settings : bpy.props.BoolProperty(default=False)
    display_bb_settings : bpy.props.BoolProperty(default=False)
    mantis_node_class_name=bl_idname[:-4]
    socket_templates = xFormBoneSockets
    
    def draw_label(self): # this will prefer a user-set label, or return the evaluated name
        return main_draw_label(self)
        
    def display_update(self, parsed_tree, context):
        # let's setup the outputs for the custom properties:
        # we'll start by trying to do it without the special mantis data
        if self.id_data.mantis_version[1] < 13: return #or the old custom properties will get wiped.
        custom_prop_display_update(self)
        if context.space_data:
            mantis_node = parsed_tree.get(get_signature_from_edited_tree(self, context))
            if mantis_node and mantis_node.prepared:
                custom_prop_display_update(self, mantis_node)
            self.display_ik_settings = False
            if mantis_node and (pb := mantis_node.bGetObject(mode='POSE')):
                self.display_ik_settings = pb.is_in_ik_chain
            self.inputs['Name'].display_text = ""
            if mantis_node and mantis_node.prepared:
                try:
                    self.inputs['Name'].display_text = mantis_node.evaluate_input("Name")
                    self.display_vp_settings = mantis_node.inputs["Custom Object"].is_connected
                    self.display_def_settings = mantis_node.evaluate_input("Deform")
                    self.display_bb_settings = mantis_node.evaluate_input("BBone Segments") > 1
                except KeyError:
                    return # the tree isn't ready yet.
            
            for template in self.socket_templates:
                name = template.name
                category = template.category
                if category == "IK":
                    self.inputs[name].hide = not self.display_ik_settings
                if category == "Display":
                    if name in ['Custom Object', 'Bone Collection']: continue
                    self.inputs[name].hide = not self.display_vp_settings
                if category == "Deform":
                    if name in ['Deform']: continue
                    self.inputs[name].hide = not self.display_def_settings
                if category == "bbone":
                    if name in ['BBone Segments']: continue
                    self.inputs[name].hide = not self.display_bb_settings

def object_displaty_update(self, parsed_tree, context):
    custom_prop_display_update(self)
    if context.space_data:
        mantis_node = parsed_tree.get(get_signature_from_edited_tree(self, context))
        self.inputs['Name'].display_text = ""
        if mantis_node and mantis_node.prepared:
            self.inputs['Name'].display_text = mantis_node.evaluate_input("Name")
            custom_prop_display_update(self, mantis_node)

class xFormArmatureNode(Node, xFormNode):
    '''A node representing an Armature object node'''
    bl_idname = 'xFormArmatureNode'
    bl_label = "Armature"
    bl_icon = 'OUTLINER_OB_ARMATURE'
    mantis_node_class_name=bl_idname[:-4]
    socket_templates = xFormArmatureSockets

    def draw_label(self): # this will prefer a user-set label, or return the evaluated name
        return main_draw_label(self)
    
    def display_update(self, parsed_tree, context):
        object_displaty_update(self, parsed_tree, context)

class xFormGeometryObjectNode(Node, xFormNode):
    """Represents a curve or mesh object."""
    bl_idname = "xFormGeometryObject"
    bl_label = "Geometry Object"
    bl_icon = "EMPTY_AXIS"
    mantis_node_class_name=bl_idname
    socket_templates = xFormGeometryObjectSockets

    def draw_label(self): # this will prefer a user-set label, or return the evaluated name
        return main_draw_label(self)
    
    def display_update(self, parsed_tree, context):
        object_displaty_update(self, parsed_tree, context)

class xFormObjectInstance(Node, xFormNode):
    """Represents an instance of an existing geometry object."""
    bl_idname = "xFormObjectInstance"
    bl_label = "Object Instance"
    bl_icon = "EMPTY_AXIS"
    mantis_node_class_name=bl_idname
    socket_templates = xFormGeometryObjectInstanceSockets

    def draw_label(self): # this will prefer a user-set label, or return the evaluated name
        return main_draw_label(self)
    
    def display_update(self, parsed_tree, context):
        object_displaty_update(self, parsed_tree, context)

class xFormCurvePin(Node, xFormNode):
    """"A node representing a curve pin"""
    bl_idname = "xFormCurvePin"
    bl_label = "Curve Pin"
    bl_icon = "FORCE_CURVE"
    mantis_node_class_name=bl_idname
    socket_templates = xFormCurvePinSockets

class xFormGetBone(Node, xFormNode):
    """"A node representing an existing bone"""
    bl_idname = "xFormGetBone"
    bl_label = "Existing Bone"
    bl_icon = "BONE_DATA"
    armature : bpy.props.StringProperty(default="")
    mantis_node_class_name=bl_idname
    socket_templates = xFormGetBoneSockets

    def display_update(self, parsed_tree, context):
        if context.space_data:
            mantis_node = parsed_tree.get(get_signature_from_edited_tree(self, context))
            if mantis_node:
                # need this parent armature to use the bone picker.
                self.inputs["Bone"].armature = mantis_node.bGetParentArmature()