LNXSDK/leenkx/blender/lnx/logicnode/lnx_node_group.py
2025-01-22 16:18:30 +01:00

578 lines
23 KiB
Python

# Some parts of this code is reused from project Sverchok.
# https://github.com/nortikin/sverchok/blob/master/core/node_group.py
#
# SPDX-License-Identifier: GPL3
# License-Filename: LICENSE
from functools import reduce
from typing import Iterator, List, Set, Dict
import bpy
from bpy.props import *
import bpy.types
from mathutils import Vector
import lnx
import lnx.logicnode.lnx_nodes as lnx_nodes
from lnx.logicnode.lnx_nodes import LnxLogicTreeNode
import lnx.utils
import lnx.props_ui
if lnx.is_reload(__name__):
lnx_nodes = lnx.reload_module(lnx_nodes)
from lnx.logicnode.lnx_nodes import LnxLogicTreeNode
lnx.utils = lnx.reload_module(lnx.utils)
lnx.props_ui = lnx.reload_module(lnx.props_ui)
else:
lnx.enable_reload(__name__)
array_nodes = lnx.logicnode.lnx_nodes.array_nodes
class LnxGroupTree(bpy.types.NodeTree):
"""Separate tree class for subtrees"""
bl_idname = 'LnxGroupTree'
bl_icon = 'NODETREE'
bl_label = 'Group Tree'
# should be updated by "Go to edit group tree" operator
group_node_name: bpy.props.StringProperty(options={'SKIP_SAVE'})
@classmethod
def poll(cls, context):
return False # only for internal usage
def upstream_trees(self) -> List['LnxGroupTree']:
"""
Returns all the tree subtrees (in case if there are group nodes)
and subtrees of subtrees and so on
The method can help predict if linking new subtree can lead to cyclic linking
"""
next_group_nodes = [node for node in self.nodes if node.bl_idname == 'LNCallGroupNode']
trees = [self]
safe_counter = 0
while next_group_nodes:
next_node = next_group_nodes.pop()
if next_node.group_tree:
trees.append(next_node.group_tree)
next_group_nodes.extend([
node for node in next_node.group_tree.nodes if node.bl_idname == 'LNCallGroupNode'])
safe_counter += 1
if safe_counter > 1000:
raise RecursionError(f'Looks like group tree "{self}" has links to itself from other groups')
return trees
def can_be_linked(self) -> bool:
"""Try to avoid creating loops of group trees with each other"""
# upstream trees of tested treed should nad share trees with downstream trees of current tree
tested_tree_upstream_trees = {t.name for t in self.upstream_trees()}
current_tree_downstream_trees = {p.node_tree.name for p in bpy.context.space_data.path}
shared_trees = tested_tree_upstream_trees & current_tree_downstream_trees
return not shared_trees
def update(self):
pass
@classmethod
def get_linkable_group_trees(cls) -> Iterator['LnxGroupTree']:
return filter(lambda tree: isinstance(tree, LnxGroupTree) and tree.can_be_linked(), bpy.data.node_groups)
@classmethod
def has_linkable_group_trees(cls) -> bool:
try:
_ = next(cls.get_linkable_group_trees())
except StopIteration:
return False
return True
class LnxEditGroupTree(bpy.types.Operator):
"""Go into subtree to edit"""
bl_idname = 'lnx.edit_group_tree'
bl_label = 'Edit Group Tree'
node_index: StringProperty(name='Node Index', default='')
def custom_poll(self, context):
if not self.node_index == '':
return True
if context.space_data.type == 'NODE_EDITOR':
if context.active_node and hasattr(context.active_node, 'group_tree'):
if context.active_node.group_tree is not None:
return True
return False
def execute(self, context):
if self.custom_poll(context):
global array_nodes
if not self.node_index == '':
group_node = array_nodes[self.node_index]
else:
group_node = context.active_node
sub_tree: LnxLogicTree = group_node.group_tree
context.space_data.path.append(sub_tree, node=group_node)
sub_tree.group_node_name = group_node.name
self.node_index = ''
return {'FINISHED'}
return {'CANCELLED'}
class LnxCopyGroupTree(bpy.types.Operator):
"""Create a copy of this group tree and use it"""
bl_idname = 'lnx.copy_group_tree'
bl_label = 'Copy Group Tree'
node_index: StringProperty(name='Node Index', default='')
def execute(self, context):
global array_nodes
group_node = array_nodes[self.node_index]
group_tree = group_node.group_tree
[setattr(n, 'copy_override', True) for n in group_tree.nodes if n.bl_idname in {'LNGroupInputsNode', 'LNGroupOutputsNode'}]
new_group_tree = group_node.group_tree.copy()
[setattr(n, 'copy_override', False) for n in group_tree.nodes if n.bl_idname in {'LNGroupInputsNode', 'LNGroupOutputsNode'}]
group_node.group_tree = new_group_tree
return {'FINISHED'}
class LnxUnlinkGroupTree(bpy.types.Operator):
"""Unlink node-group (Shift + Click to set users to zero, data will then not be saved)"""
bl_idname = 'lnx.unlink_group_tree'
bl_label = 'Unlink Group Tree'
node_index: StringProperty(name='Node Index', default='')
def invoke(self, context, event):
self.clear = False
if event.shift:
self.clear = True
self.execute(context)
return {'FINISHED'}
def execute(self, context):
global array_nodes
group_node = array_nodes[self.node_index]
group_tree = group_node.group_tree
group_node.group_tree = None
if self.clear:
group_tree.user_clear()
return {'FINISHED'}
class LnxSearchGroupTree(bpy.types.Operator):
"""Browse group trees to be linked"""
bl_idname = 'lnx.search_group_tree'
bl_label = 'Search Group Tree'
bl_property = 'tree_name'
node_index: StringProperty(name='Node Index', default='')
def available_trees(self, context):
linkable_trees = filter(lambda t: hasattr(t, 'can_be_linked') and t.can_be_linked(), bpy.data.node_groups)
return [(t.name, ('0 ' if t.users == 0 else 'F ' if t.use_fake_user else '') + t.name, '') for t in linkable_trees]
tree_name: bpy.props.EnumProperty(items=available_trees)
def execute(self, context):
global array_nodes
tree_to_link = bpy.data.node_groups[self.tree_name]
group_node = array_nodes[self.node_index]
group_node.group_tree = tree_to_link
return {'FINISHED'}
def invoke(self, context, event):
context.window_manager.invoke_search_popup(self)
return {'FINISHED'}
class LnxAddGroupTree(bpy.types.Operator):
"""Create empty subtree for group node"""
bl_idname = "lnx.add_group_tree"
bl_label = "Add Group Tree"
node_index: StringProperty(name='Node Index', default='')
@classmethod
def poll(cls, context):
path = getattr(context.space_data, 'path', [])
if len(path):
if path[-1].node_tree.bl_idname in {'LnxLogicTreeType', 'LnxGroupTree'}:
return True
return False
def execute(self, context):
"""Link new subtree to group node, create input and output nodes in subtree and go to edit one"""
global array_nodes
sub_tree = bpy.data.node_groups.new('Leenkx group', 'LnxGroupTree') # creating subtree
sub_tree.use_fake_user = True
group_node = array_nodes[self.node_index]
group_node.group_tree = sub_tree # link subtree to group node
sub_tree.nodes.new('LNGroupInputsNode').location = (-250, 0) # create node for putting data into subtree
sub_tree.nodes.new('LNGroupOutputsNode').location = (250, 0) # create node for getting data from subtree
context.space_data.path.append(sub_tree, node=group_node)
sub_tree.group_node_name = group_node.name
return {'FINISHED'}
class LnxAddGroupTreeFromSelected(bpy.types.Operator):
"""Select nodes group node and placing them into subtree"""
bl_idname = "lnx.add_group_tree_from_selected"
bl_label = "Create Group Tree from Selected"
@classmethod
def poll(cls, context):
path = getattr(context.space_data, 'path', [])
if len(path):
if path[-1].node_tree.bl_idname in {'LnxLogicTreeType', 'LnxGroupTree'}:
return bool(cls.filter_selected_nodes(path[-1].node_tree))
return False
def execute(self, context):
"""
Add group tree from selected:
01. Deselect group Input and Output nodes
02. Copy nodes into clipboard
03. Create group tree and move into one
04. Past nodes from clipboard
05. Move nodes into tree center
06. Add group "input" and "output" outside of bounding box of the nodes
07. TODO: Connect "input" and "output" sockets with group nodes
08. Add Group tree node in center of selected node in initial tree
09. TODO: Link the node with appropriate sockets
10. Cleaning
"""
base_tree = context.space_data.path[-1].node_tree
sub_tree: LnxGroupTree = bpy.data.node_groups.new('Leenkx group', 'LnxGroupTree')
# deselect group nodes if selected
[setattr(n, 'select', False) for n in base_tree.nodes if n.select and n.bl_idname in {'LNGroupInputsNode', 'LNGroupOutputsNode'}]
# Frames can't be just copied because they do not have absolute location, but they can be recreated
frame_names = {n.name for n in base_tree.nodes if n.select and n.bl_idname == 'NodeFrame'}
[setattr(n, 'select', False) for n in base_tree.nodes if n.bl_idname == 'NodeFrame']
# copy and past nodes into group tree
bpy.ops.node.clipboard_copy()
context.space_data.path.append(sub_tree)
bpy.ops.node.clipboard_paste()
context.space_data.path.pop() # will enter later via operator
# move nodes in tree center
sub_tree_nodes = self.filter_selected_nodes(sub_tree)
center = reduce(lambda v1, v2: v1 + v2, [n.location for n in sub_tree_nodes]) / len(sub_tree_nodes)
[setattr(n, 'location', n.location - center) for n in sub_tree_nodes]
# recreate frames
node_name_mapping = {n.name: n.name for n in sub_tree.nodes} # all nodes have the same name as in base tree
self.recreate_frames(base_tree, sub_tree, frame_names, node_name_mapping)
# add group input and output nodes
min_x = min(n.location[0] for n in sub_tree_nodes)
max_x = max(n.location[0] for n in sub_tree_nodes)
input_node = sub_tree.nodes.new('LNGroupInputsNode')
input_node.location = (min_x - 250, 0)
output_node = sub_tree.nodes.new('LNGroupOutputsNode')
output_node.location = (max_x + 250, 0)
# add group tree node
initial_nodes = self.filter_selected_nodes(base_tree)
center = reduce(lambda v1, v2: v1 + v2,
[Vector(LnxLogicTreeNode.absolute_location(n)) for n in initial_nodes]) / len(initial_nodes)
group_node = base_tree.nodes.new('LNCallGroupNode')
group_node.select = False
group_node.group_tree = sub_tree
group_node.location = center
sub_tree.group_node_name = group_node.name
# delete selected nodes and copied frames without children
[base_tree.nodes.remove(n) for n in self.filter_selected_nodes(base_tree)]
with_children_frames = {n.parent.name for n in base_tree.nodes if n.parent}
[base_tree.nodes.remove(n) for n in base_tree.nodes if n.name in frame_names and n.name not in with_children_frames]
# enter the group tree
bpy.ops.lnx.edit_group_tree(node_index=group_node.get_id_str())
return {'FINISHED'}
@staticmethod
def filter_selected_nodes(tree) -> list:
"""Avoiding selecting nodes which should not be copied into subtree"""
return [n for n in tree.nodes if n.select and n.bl_idname not in {'LNGroupInputsNode', 'LNGroupOutputsNode'}]
@staticmethod
def recreate_frames(from_tree: bpy.types.NodeTree,
to_tree: bpy.types.NodeTree,
frame_names: Set[str],
from_to_node_names: Dict[str, str]):
"""
Copy frames from one tree to another
from_to_node_names - mapping of node names between two trees
"""
new_frame_names = {n: to_tree.nodes.new('NodeFrame').name for n in frame_names}
frame_attributes = ['label', 'use_custom_color', 'color', 'label_size', 'text']
for frame_name in frame_names:
old_frame = from_tree.nodes[frame_name]
new_frame = to_tree.nodes[new_frame_names[frame_name]]
for attr in frame_attributes:
setattr(new_frame, attr, getattr(old_frame, attr))
for from_node in from_tree.nodes:
if from_node.name not in from_to_node_names:
continue
if from_node.parent and from_node.parent.name in new_frame_names:
if from_node.bl_idname == 'NodeFrame':
to_node = to_tree.nodes[new_frame_names[from_node.name]]
else:
to_node = to_tree.nodes[from_to_node_names[from_node.name]]
to_node.parent = to_tree.nodes[new_frame_names[from_node.parent.name]]
class TreeVarNameConflictItem(bpy.types.PropertyGroup):
"""Represents two conflicting tree variables with the same name"""
name: StringProperty(
description='The name of the conflicting tree variables'
)
action: EnumProperty(
name='Conflict Resolution Action',
description='How to resolve the tree variable conflict',
default='rename',
items=[
('rename', 'Rename', 'Automatically rename the group\'s tree variable'),
('merge', 'Merge', 'Merge the conflicting tree variables'),
]
)
needs_rename: BoolProperty(
description='If true, the conflict needs to be resolved by renaming'
)
class LnxUngroupGroupTree(bpy.types.Operator):
"""Put sub nodes into current layout and delete current group node"""
bl_idname = 'lnx.ungroup_group_tree'
bl_label = "Ungroup Group Tree"
bl_options = {'UNDO'} # Required to "un-rename" node's lnx_logic_id in case of tree variable conflicts
conflicts: CollectionProperty(type=TreeVarNameConflictItem)
@classmethod
def poll(cls, context):
if context.space_data.type == 'NODE_EDITOR':
if context.active_node and hasattr(context.active_node, 'group_tree'):
if context.active_node.group_tree is not None:
return True
return False
def invoke(self, context, event):
group_node = context.active_node
group_tree = group_node.group_tree
dest_tree = group_node.get_tree()
# name -> type
group_tree_variables = {}
dest_tree_variables = {}
for var in group_tree.lnx_treevariableslist:
group_tree_variables[var.name] = var.node_type
for var in dest_tree.lnx_treevariableslist:
dest_tree_variables[var.name] = var.node_type
# Check for conflicting tree variables
self.conflicts.clear() # Might still contain values from previous invocation
conflicting_var_names = group_tree_variables.keys() & dest_tree_variables.keys()
user_can_choose = False
for conflicting_var_name in conflicting_var_names:
conflict_item = self.conflicts.add()
conflict_item.name = conflicting_var_name
# Tree variable types differ, cannot merge
conflict_item.needs_rename = group_tree_variables[conflicting_var_name] != dest_tree_variables[conflicting_var_name]
user_can_choose |= not conflict_item.needs_rename
# If there are no conflicts or all conflicts _must_ be resolved
# via renaming there's no reason to ask the user
if user_can_choose:
wm = context.window_manager
return wm.invoke_props_dialog(self, width=400)
return self.execute(context)
def draw(self, context):
layout = self.layout
lnx.props_ui.draw_multiline_with_icon(
layout, layout_width_px=400,
icon='ERROR',
text=(
'The group\'s logic tree contains tree variables whose names'
' are identical to tree variable names in the enclosing tree.'
)
)
layout.label(icon='BLANK1', text='Please choose how to resolve the naming conflicts (press ESC to cancel):')
layout.separator()
conflict_item: TreeVarNameConflictItem
for conflict_item in self.conflicts:
split = layout.split(factor=0.6)
split.alignment = 'RIGHT'
split.label(text=conflict_item.name)
if conflict_item.needs_rename:
row = split.row()
row.label(text="Needs rename")
else:
row = split.row()
row.prop(conflict_item, "action", expand=True)
layout.separator() # Add space above Blender's "OK" button
def execute(self, context):
"""Similar to AddGroupTreeFromSelected operator but in backward direction (from subtree to tree)"""
# go to subtree, select all except input and output groups and mark nodes to be copied
group_node = context.active_node
sub_tree = group_node.group_tree
if len(self.conflicts) > 0:
self._resolve_conflicts(sub_tree, group_node.get_tree())
bpy.ops.lnx.edit_group_tree(node_index=group_node.get_id_str())
[setattr(n, 'select', False) for n in sub_tree.nodes]
group_nodes_filter = filter(lambda n: n.bl_idname not in {'LNGroupInputsNode', 'LNGroupOutputsNode'}, sub_tree.nodes)
for node in group_nodes_filter:
node.select = True
node['sub_node_name'] = node.name # this will be copied within the nodes
# the attribute should be empty in destination tree
tree = context.space_data.path[-2].node_tree
for node in tree.nodes:
if 'sub_node_name' in node:
del node['sub_node_name']
# Frames can't be just copied because they do not have absolute location, but they can be recreated
frame_names = {n.name for n in sub_tree.nodes if n.select and n.bl_idname == 'NodeFrame'}
[setattr(n, 'select', False) for n in sub_tree.nodes if n.bl_idname == 'NodeFrame']
if any(n for n in sub_tree.nodes if n.select): # if no selection copy operator will raise error
# copy and past nodes into group tree
bpy.ops.node.clipboard_copy()
context.space_data.path.pop()
bpy.ops.node.clipboard_paste() # this will deselect all and select only pasted nodes
# move nodes in group node center
tree_select_nodes = [n for n in tree.nodes if n.select]
center = reduce(lambda v1, v2: v1 + v2,
[Vector(LnxLogicTreeNode.absolute_location(n)) for n in tree_select_nodes]) / len(tree_select_nodes)
[setattr(n, 'location', n.location - (center - group_node.location)) for n in tree_select_nodes]
# recreate frames
node_name_mapping = {n['sub_node_name']: n.name for n in tree.nodes if 'sub_node_name' in n}
LnxAddGroupTreeFromSelected.recreate_frames(sub_tree, tree, frame_names, node_name_mapping)
else:
context.space_data.path.pop() # should exit from subtree anywhere
# delete group node
tree.nodes.remove(group_node)
for node in tree.nodes:
if 'sub_node_name' in node:
del node['sub_node_name']
tree.update()
return {'FINISHED'}
def _resolve_conflicts(self, group_tree: bpy.types.NodeTree, dest_tree: bpy.types.NodeTree):
rename_conflict_names = {} # old variable name -> new variable name
for conflict_item in self.conflicts:
if conflict_item.needs_rename or conflict_item.action == 'rename':
# Initialize as empty, will be set further below
rename_conflict_names[conflict_item.name] = ''
for var_item in group_tree.lnx_treevariableslist:
if var_item.name in rename_conflict_names:
# Create a renamed variable in the destination tree and ensure
# its name doesn't conflict with either tree
new_name = group_tree.name + '.' + var_item.name
dest_var = dest_tree.lnx_treevariableslist.add()
dest_varname = lnx.utils.unique_name_in_lists(
item_lists=[group_tree.lnx_treevariableslist, dest_tree.lnx_treevariableslist],
name_attr='name', wanted_name=new_name, ignore_item=dest_var
)
dest_var['_name'] = dest_varname
rename_conflict_names[var_item.name] = dest_varname
dest_var.node_type = var_item.node_type
dest_var.color = var_item.color
# Update the logic ids so that copying the nodes to the new tree
# pastes them with references to the newly created dest_var
for node in group_tree.nodes:
node.lnx_logic_id = rename_conflict_names.get(node.lnx_logic_id, node.lnx_logic_id)
class LnxAddCallGroupNode(bpy.types.Operator):
"""Create A Call Group Node"""
bl_idname = 'lnx.add_call_group_node'
bl_label = "Add 'Call Node Group' Node"
node_ref = None
@classmethod
def poll(cls, context):
if context.space_data.type == 'NODE_EDITOR':
return context.space_data.edit_tree and context.space_data.tree_type == 'LnxLogicTreeType'
return False
def invoke(self, context, event):
context.window_manager.modal_handler_add(self)
self.execute(context)
return {'RUNNING_MODAL'}
def modal(self, context, event):
if event.type == 'MOUSEMOVE':
self.node_ref.location = context.space_data.cursor_location
elif event.type == 'LEFTMOUSE': # Confirm
return {'FINISHED'}
return {'RUNNING_MODAL'}
def execute(self, context):
tree = context.space_data.path[-1].node_tree
self.node_ref = tree.nodes.new('LNCallGroupNode')
self.node_ref.location = context.space_data.cursor_location
return {'FINISHED'}
class LNX_PT_LogicGroupPanel(bpy.types.Panel):
bl_label = 'Leenkx Logic Group'
bl_idname = 'LNX_PT_LogicGroupPanel'
bl_space_type = 'NODE_EDITOR'
bl_region_type = 'UI'
bl_category = 'Leenkx'
@classmethod
def poll(cls, context):
return context.space_data.tree_type == 'LnxLogicTreeType' and context.space_data.edit_tree
def has_active_node(self, context):
if context.active_node and hasattr(context.active_node, 'group_tree'):
if context.active_node.group_tree is not None:
return True
return False
def draw(self, context):
layout = self.layout
layout.operator('lnx.add_call_group_node', icon='ADD')
layout.operator('lnx.add_group_tree_from_selected', icon='NODETREE')
layout.operator('lnx.ungroup_group_tree', icon='NODETREE')
row = layout.row()
row.enabled = self.has_active_node(context)
row.operator('lnx.edit_group_tree', icon='FULLSCREEN_ENTER', text='Edit Tree')
__REG_CLASSES = (
LnxGroupTree,
LnxEditGroupTree,
LnxCopyGroupTree,
LnxUnlinkGroupTree,
LnxSearchGroupTree,
LnxAddGroupTree,
LnxAddGroupTreeFromSelected,
TreeVarNameConflictItem,
LnxUngroupGroupTree,
LnxAddCallGroupNode,
LNX_PT_LogicGroupPanel
)
register, unregister = bpy.utils.register_classes_factory(__REG_CLASSES)