armory/blender/arm/logicnode/arm_nodes.py

import itertools
from collections import OrderedDict
from typing import Any, Generator, List, Optional, Type, Dict
from typing import OrderedDict as ODict  # Prevent naming conflicts

import bpy.types
from bpy.props import *
from nodeitems_utils import NodeItem
from arm.logicnode.arm_sockets import ArmCustomSocket

# Pass NodeReplacment forward to individual node modules that import arm_nodes
from arm.logicnode.replacement import NodeReplacement
import arm.node_utils

# When passed as a category to add_node(), this will use the capitalized
# name of the package of the node as the category to make renaming
# categories easier.
PKG_AS_CATEGORY = "__pkgcat__"

nodes = []
category_items: ODict[str, List['ArmNodeCategory']] = OrderedDict()

array_nodes = dict()


class ArmLogicTreeNode(bpy.types.Node):
    arm_category = PKG_AS_CATEGORY
    arm_section = 'default'
    arm_is_obsolete = False

    def init(self, context):
        # make sure a given node knows the version of the NodeClass from when it was created
        if isinstance(type(self).arm_version, int):
            self.arm_version = type(self).arm_version
        else:
            self.arm_version = 1

    @classmethod
    def poll(cls, ntree):
        return ntree.bl_idname == 'ArmLogicTreeType'

    @classmethod
    def on_register(cls):
        """Don't call this method register() as it will be triggered before Blender registers the class, resulting in
        a double registration."""
        add_node(cls, cls.arm_category, cls.arm_section, cls.arm_is_obsolete)

    @classmethod
    def on_unregister(cls):
        pass

    def get_replacement_node(self, node_tree: bpy.types.NodeTree):
        # needs to be overridden by individual node classes with arm_version>1
        """(only called if the node's version is inferior to the node class's version)
        Help with the node replacement process, by explaining how a node (`self`) should be replaced.
        This method can either return a NodeReplacement object (see `nodes_logic.py`), or a brand new node.

        If a new node is returned, then the following needs to be already set:
        - the node's links to the other nodes
        - the node's properties
        - the node inputs's default values

        If more than one node need to be created (for example, if an input needs a type conversion after update),
        please return all the nodes in a list.

        please raise a LookupError specifically when the node's version isn't handled by the function.

        note that the lowest 'defined' version should be 1. if the node's version is 0, it means that it has been saved before versioning was a thing.
        NODES OF VERSION 1 AND VERSION 0 SHOULD HAVE THE SAME CONTENTS
        """
        if self.arm_version == 0 and type(self).arm_version == 1:
            # In case someone doesn't implement this function, but the node has version 0
            return NodeReplacement.Identity(self)
        else:
            raise LookupError(f"the current node class {repr(type(self)):s} does not implement get_replacement_node() even though it has updated")

    def add_input(self, socket_type: str, socket_name: str, default_value: Any = None, is_var: bool = False) -> bpy.types.NodeSocket:
        """Adds a new input socket to the node.

        If `is_var` is true, a dot is placed inside the socket to denote
        that this socket can be used for variable access (see
        SetVariable node).
        """
        socket = self.inputs.new(socket_type, socket_name)

        if default_value is not None:
            if isinstance(socket, ArmCustomSocket):
                if socket.arm_socket_type != 'NONE':
                    socket.default_value_raw = default_value
                else:
                    raise ValueError('specified a default value for an input node that doesn\'t accept one')
            else:
                socket.default_value = default_value

        if is_var and not socket.display_shape.endswith('_DOT'):
            socket.display_shape += '_DOT'

        return socket

    def add_output(self, socket_type: str, socket_name: str, default_value: Any = None, is_var: bool = False) -> bpy.types.NodeSocket:
        """Adds a new output socket to the node.

        If `is_var` is true, a dot is placed inside the socket to denote
        that this socket can be used for variable access (see
        SetVariable node).
        """
        socket = self.outputs.new(socket_type, socket_name)

        if default_value is not None:
            socket.default_value = default_value

        if is_var and not socket.display_shape.endswith('_DOT'):
            socket.display_shape += '_DOT'

        return socket


class ArmNodeAddInputButton(bpy.types.Operator):
    """Add a new input socket to the node set by node_index."""
    bl_idname = 'arm.node_add_input'
    bl_label = 'Add Input'

    node_index: StringProperty(name='Node Index', default='')
    socket_type: StringProperty(name='Socket Type', default='NodeSocketShader')
    name_format: StringProperty(name='Name Format', default='Input {0}')
    index_name_offset: IntProperty(name='Index Name Offset', default=0)

    def execute(self, context):
        global array_nodes
        inps = array_nodes[self.node_index].inputs

        socket_types = self.socket_type.split(';')
        name_formats = self.name_format.split(';')
        assert len(socket_types)==len(name_formats)
        
        format_index = (len(inps) + self.index_name_offset) //len(socket_types)
        for socket_type, name_format in zip(socket_types, name_formats):
            inps.new(socket_type, name_format.format(str(format_index)))

        # Reset to default again for subsequent calls of this operator
        self.node_index = ''
        self.socket_type = 'NodeSocketShader'
        self.name_format = 'Input {0}'
        self.index_name_offset = 0

        return{'FINISHED'}

class ArmNodeAddInputValueButton(bpy.types.Operator):
     """Add new input"""
     bl_idname = 'arm.node_add_input_value'
     bl_label = 'Add Input'
     node_index: StringProperty(name='Node Index', default='')
     socket_type: StringProperty(name='Socket Type', default='NodeSocketShader')

     def execute(self, context):
         global array_nodes
         inps = array_nodes[self.node_index].inputs
         inps.new(self.socket_type, 'Value')
         return{'FINISHED'}

class ArmNodeRemoveInputButton(bpy.types.Operator):
    """Remove last input"""
    bl_idname = 'arm.node_remove_input'
    bl_label = 'Remove Input'
    node_index: StringProperty(name='Node Index', default='')
    count: IntProperty(name='Number of inputs to remove', default=1, min=1)
    min_inputs: IntProperty(name='Number of inputs to keep', default=0, min=0)

    def execute(self, context):
        global array_nodes
        node = array_nodes[self.node_index]
        inps = node.inputs
        min_inps = self.min_inputs if not hasattr(node, 'min_inputs') else node.min_inputs
        if len(inps) >= min_inps + self.count:
            for _ in range(self.count):
                inps.remove(inps.values()[-1])
        return{'FINISHED'}

class ArmNodeRemoveInputValueButton(bpy.types.Operator):
    """Remove last input"""
    bl_idname = 'arm.node_remove_input_value'
    bl_label = 'Remove Input'
    node_index: StringProperty(name='Node Index', default='')
    target_name: StringProperty(name='Name of socket to remove', default='Value')

    def execute(self, context):
        global array_nodes
        node = array_nodes[self.node_index]
        inps = node.inputs
        min_inps = 0 if not hasattr(node, 'min_inputs') else node.min_inputs
        if len(inps) > min_inps and inps[-1].name == self.target_name:
            inps.remove(inps.values()[-1])
        return{'FINISHED'}

class ArmNodeAddOutputButton(bpy.types.Operator):
    """Add a new output socket to the node set by node_index"""
    bl_idname = 'arm.node_add_output'
    bl_label = 'Add Output'

    node_index: StringProperty(name='Node Index', default='')
    socket_type: StringProperty(name='Socket Type', default='NodeSocketShader')
    name_format: StringProperty(name='Name Format', default='Output {0}')
    index_name_offset: IntProperty(name='Index Name Offset', default=0)

    def execute(self, context):
        global array_nodes
        outs = array_nodes[self.node_index].outputs

        socket_types = self.socket_type.split(';')
        name_formats = self.name_format.split(';')
        assert len(socket_types)==len(name_formats)
        
        format_index = (len(outs) + self.index_name_offset) //len(socket_types)
        for socket_type, name_format in zip(socket_types, name_formats):
            outs.new(socket_type, name_format.format(str(format_index)))

        # Reset to default again for subsequent calls of this operator
        self.node_index = ''
        self.socket_type = 'NodeSocketShader'
        self.name_format = 'Output {0}'
        self.index_name_offset = 0

        return{'FINISHED'}

class ArmNodeRemoveOutputButton(bpy.types.Operator):
    """Remove last output"""
    bl_idname = 'arm.node_remove_output'
    bl_label = 'Remove Output'
    node_index: StringProperty(name='Node Index', default='')
    count: IntProperty(name='Number of outputs to remove', default=1, min=1)

    def execute(self, context):
        global array_nodes
        node = array_nodes[self.node_index]
        outs = node.outputs
        min_outs = 0 if not hasattr(node, 'min_outputs') else node.min_outputs
        if len(outs) >= min_outs + self.count:
            for _ in range(self.count):
                outs.remove(outs.values()[-1])            
        return{'FINISHED'}

class ArmNodeAddInputOutputButton(bpy.types.Operator):
    """Add new input and output"""
    bl_idname = 'arm.node_add_input_output'
    bl_label = 'Add Input Output'

    node_index: StringProperty(name='Node Index', default='')
    in_socket_type: StringProperty(name='In Socket Type', default='NodeSocketShader')
    out_socket_type: StringProperty(name='Out Socket Type', default='NodeSocketShader')
    in_name_format: StringProperty(name='In Name Format', default='Input {0}')
    out_name_format: StringProperty(name='Out Name Format', default='Output {0}')
    in_index_name_offset: IntProperty(name='Index Name Offset', default=0)

    def execute(self, context):
        global array_nodes
        node = array_nodes[self.node_index]
        inps = node.inputs
        outs = node.outputs

        in_socket_types = self.in_socket_type.split(';')
        in_name_formats = self.in_name_format.split(';')
        assert len(in_socket_types)==len(in_name_formats)

        out_socket_types = self.out_socket_type.split(';')
        out_name_formats = self.out_name_format.split(';')
        assert len(out_socket_types)==len(out_name_formats)
        
        in_format_index = (len(outs) + self.index_name_offset) // len(in_socket_types)
        out_format_index = len(outs) // len(out_socket_types)
        for socket_type, name_format in zip(in_socket_types, in_name_formats):
            inps.new(socket_type, name_format.format(str(in_format_index)))
        for socket_type, name_format in zip(out_socket_types, out_name_formats):
            outs.new(socket_type, name_format.format(str(out_format_index)))

        # Reset to default again for subsequent calls of this operator
        self.node_index = ''
        self.in_socket_type = 'NodeSocketShader'
        self.out_socket_type = 'NodeSocketShader'
        self.in_name_format = 'Input {0}'
        self.out_name_format = 'Output {0}'
        self.in_index_name_offset = 0

        return{'FINISHED'}

class ArmNodeRemoveInputOutputButton(bpy.types.Operator):
    """Remove last input and output"""
    bl_idname = 'arm.node_remove_input_output'
    bl_label = 'Remove Input Output'
    node_index: StringProperty(name='Node Index', default='')
    in_count: IntProperty(name='Number of inputs to remove', default=1, min=1)
    out_count: IntProperty(name='Number of inputs to remove', default=1, min=1)
    
    def execute(self, context):
        global array_nodes
        node = array_nodes[self.node_index]
        inps = node.inputs
        outs = node.outputs
        min_inps = 0 if not hasattr(node, 'min_inputs') else node.min_inputs
        min_outs = 0 if not hasattr(node, 'min_outputs') else node.min_outputs
        if len(inps) >= min_inps + self.in_count:
            for _ in range(self.in_count):
                inps.remove(inps.values()[-1])
        if len(outs) >= min_outs + self.out_count:
            for _ in range(self.out_count):
                outs.remove(outs.values()[-1])
        return{'FINISHED'}


class ArmNodeSearch(bpy.types.Operator):
    bl_idname = "arm.node_search"
    bl_label = "Search..."
    bl_options = {"REGISTER"}
    bl_property = "item"

    def get_search_items(self, context):
        items = []
        for node in get_all_nodes():
            items.append((node.nodetype, node.label, ""))
        return items

    item: EnumProperty(items=get_search_items)

    @classmethod
    def poll(cls, context):
        return context.space_data.tree_type == 'ArmLogicTreeType' and context.space_data.edit_tree

    @classmethod
    def description(cls, context, properties):
        if cls.poll(context):
            return "Search for a logic node"
        else:
            return "Search for a logic node. This operator is not available" \
                   " without an active node tree"

    def invoke(self, context, event):
        context.window_manager.invoke_search_popup(self)
        return {"CANCELLED"}

    def execute(self, context):
        """Called when a node is added."""
        bpy.ops.node.add_node('INVOKE_DEFAULT', type=self.item, use_transform=True)
        return {"FINISHED"}


class ArmNodeCategory:
    """Represents a category (=directory) of logic nodes."""
    def __init__(self, name: str, icon: str, description: str):
        self.name = name
        self.icon = icon
        self.description = description
        self.node_sections: ODict[str, List[NodeItem]] = OrderedDict()
        self.deprecated_nodes: List[NodeItem] = []

    def register_node(self, node_type: Type[bpy.types.Node], node_section: str) -> None:
        """Registers a node to this category so that it will be
        displayed int the `Add node` menu."""
        self.add_node_section(node_section)
        self.node_sections[node_section].append(arm.node_utils.nodetype_to_nodeitem(node_type))

    def register_deprecated_node(self, node_type: Type[bpy.types.Node]) -> None:
        if hasattr(node_type, 'arm_is_obsolete') and node_type.arm_is_obsolete:
            self.deprecated_nodes.append(arm.node_utils.nodetype_to_nodeitem(node_type))

    def get_all_nodes(self) -> Generator[NodeItem, None, None]:
        """Returns all nodes that are registered into this category."""
        yield from itertools.chain(*self.node_sections.values())

    def add_node_section(self, name: str):
        """Adds a node section to this category."""
        if name not in self.node_sections:
            self.node_sections[name] = []

    def sort_nodes(self):
        for node_section in self.node_sections:
            self.node_sections[node_section] = sorted(self.node_sections[node_section], key=lambda item: item.label)


def category_exists(name: str) -> bool:
    for category_section in category_items:
        for c in category_items[category_section]:
            if c.name == name:
                return True

    return False


def get_category(name: str) -> Optional[ArmNodeCategory]:
    for category_section in category_items:
        for c in category_items[category_section]:
            if c.name == name:
                return c

    return None


def get_all_categories() -> Generator[ArmNodeCategory, None, None]:
    for section_categories in category_items.values():
        yield from itertools.chain(section_categories)


def get_all_nodes() -> Generator[NodeItem, None, None]:
    for category in get_all_categories():
        yield from itertools.chain(category.get_all_nodes())


def add_category_section(name: str) -> None:
    """Adds a section of categories to the node menu to group multiple
    categories visually together. The given name only acts as an ID and
    is not displayed in the user inferface."""
    global category_items
    if name not in category_items:
        category_items[name] = []


def add_node_section(name: str, category: str) -> None:
    """Adds a section of nodes to the sub menu of the given category to
    group multiple nodes visually together. The given name only acts as
    an ID and is not displayed in the user inferface."""
    node_category = get_category(category)

    if node_category is not None:
        node_category.add_node_section(name)


def add_category(category: str, section: str = 'default', icon: str = 'BLANK1', description: str = '') -> Optional[ArmNodeCategory]:
    """Adds a category of nodes to the node menu."""
    global category_items

    add_category_section(section)
    if not category_exists(category):
        node_category = ArmNodeCategory(category, icon, description)
        category_items[section].append(node_category)
        return node_category

    return None


def add_node(node_type: Type[bpy.types.Node], category: str, section: str = 'default', is_obsolete: bool = False) -> None:
    """
    Registers a node to the given category. If no section is given, the
    node is put into the default section that does always exist.

    Warning: Make sure that this function is not called multiple times per node!
    """
    global nodes

    if category == PKG_AS_CATEGORY:
        category = node_type.__module__.rsplit('.', 2)[-2].capitalize()

    nodes.append(node_type)
    node_category = get_category(category)

    if node_category is None:
        node_category = add_category(category)

    if is_obsolete:
        # We need the obsolete nodes to be registered in order to have them replaced,
        # but do not add them to the menu.
        if node_category is not None:
            # Make the deprecated nodes available for documentation purposes
            node_category.register_deprecated_node(node_type)
        return

    node_category.register_node(node_type, section)
    node_type.bl_icon = node_category.icon


def deprecated(*alternatives: str, message=""):
    """Class decorator to deprecate logic node classes. You can pass multiple string
    arguments with the names of the available alternatives as well as a message
    (keyword-param only) with further information about the deprecation."""

    def wrapper(cls: ArmLogicTreeNode) -> ArmLogicTreeNode:
        cls.bl_label += ' (Deprecated)'
        if hasattr(cls, 'bl_description'):
            cls.bl_description = f'Deprecated. {cls.bl_description}'
        else:
            cls.bl_description = 'Deprecated.'
        cls.bl_icon = 'ERROR'
        cls.arm_is_obsolete = True

        if cls.__doc__ is None:
            cls.__doc__ = ''

        if len(alternatives) > 0:
            cls.__doc__ += '\n' + f'@deprecated {",".join(alternatives)}: {message}'
        else:
            cls.__doc__ += '\n' + f'@deprecated : {message}'

        return cls

    return wrapper


def reset_globals():
    global nodes
    global category_items
    nodes = []
    category_items = OrderedDict()


bpy.utils.register_class(ArmNodeSearch)
bpy.utils.register_class(ArmNodeAddInputButton)
bpy.utils.register_class(ArmNodeAddInputValueButton)
bpy.utils.register_class(ArmNodeRemoveInputButton)
bpy.utils.register_class(ArmNodeRemoveInputValueButton)
bpy.utils.register_class(ArmNodeAddOutputButton)
bpy.utils.register_class(ArmNodeRemoveOutputButton)
bpy.utils.register_class(ArmNodeAddInputOutputButton)
bpy.utils.register_class(ArmNodeRemoveInputOutputButton)