Files
test2/scripts/startup/bl_operators/node.py

1337 lines
45 KiB
Python

# SPDX-FileCopyrightText: 2012-2023 Blender Authors
#
# SPDX-License-Identifier: GPL-2.0-or-later
from __future__ import annotations
import bpy
from bpy.types import (
FileHandler,
Operator,
PropertyGroup,
)
from bpy.props import (
BoolProperty,
CollectionProperty,
EnumProperty,
FloatVectorProperty,
StringProperty,
IntProperty,
)
from mathutils import (
Vector,
)
from bpy.app.translations import (
pgettext_tip as tip_,
pgettext_rpt as rpt_,
)
math_nodes = {
"ShaderNodeMath",
"ShaderNodeVectorMath",
"FunctionNodeIntegerMath",
"FunctionNodeBooleanMath",
"FunctionNodeBitMath",
}
switch_nodes = {
"GeometryNodeMenuSwitch",
"GeometryNodeIndexSwitch",
}
# A context manager for temporarily unparenting nodes from their frames
# This gets rid of issues with framed nodes using relative coordinates
class temporary_unframe:
def __init__(self, nodes):
self.parent_dict = {}
for node in nodes:
if node.parent is not None:
self.parent_dict[node] = node.parent
node.parent = None
def __enter__(self):
return self
def __exit__(self, _type, _value, _traceback):
for node, parent in self.parent_dict.items():
node.parent = parent
def cast_value(source, target):
source_type = source.type
target_type = target.type
value = source.default_value
def to_bool(value):
return value > 0
def single_value_to_color(value):
return Vector((value, value, value, 1.0))
def single_value_to_vector(value):
return Vector([value,] * len(target.default_value))
def color_to_float(color):
return (0.2126 * color[0]) + (0.7152 * color[1]) + (0.0722 * color[2])
def vector_to_float(vector):
return sum(vector) / len(vector)
func_map = {
('VALUE', 'INT'): int,
('VALUE', 'BOOLEAN'): to_bool,
('VALUE', 'RGBA'): single_value_to_color,
('VALUE', 'VECTOR'): single_value_to_vector,
('INT', 'VALUE'): float,
('INT', 'BOOLEAN'): to_bool,
('INT', 'RGBA'): single_value_to_color,
('INT', 'VECTOR'): single_value_to_vector,
('BOOLEAN', 'VALUE'): float,
('BOOLEAN', 'INT'): int,
('BOOLEAN', 'RGBA'): single_value_to_color,
('BOOLEAN', 'VECTOR'): single_value_to_vector,
('RGBA', 'VALUE'): color_to_float,
('RGBA', 'INT'): lambda color: int(color_to_float(color)),
('RGBA', 'BOOLEAN'): lambda color: to_bool(color_to_float(color)),
('RGBA', 'VECTOR'): lambda color: color[:len(target.default_value)],
('VECTOR', 'VALUE'): vector_to_float,
('VECTOR', 'INT'): lambda vector: int(vector_to_float(vector)),
# Even negative vectors get implicitly converted to True, hence to_bool is not used
('VECTOR', 'BOOLEAN'): lambda vector: bool(vector_to_float(vector)),
('VECTOR', 'RGBA'): lambda vector: list(vector).extend([0.0] * (len(target.default_value) - len(vector)))
}
if source_type == target_type:
return value
cast_func = func_map.get((source_type, target_type))
if cast_func is not None:
return cast_func(value)
return None
class NodeSetting(PropertyGroup):
__slots__ = ()
value: StringProperty(
name="Value",
description="Python expression to be evaluated "
"as the initial node setting",
default="",
)
class NodeOperator:
settings: CollectionProperty(
name="Settings",
description="Settings to be applied on the newly created node",
type=NodeSetting,
options={'SKIP_SAVE'},
)
@classmethod
def description(cls, _context, properties):
from nodeitems_builtins import node_tree_group_type
nodetype = properties["type"]
if nodetype in node_tree_group_type.values():
for setting in properties.settings:
if setting.name == "node_tree":
node_group = eval(setting.value)
if node_group.description:
return node_group.description
bl_rna = bpy.types.Node.bl_rna_get_subclass(nodetype)
if bl_rna is not None:
return tip_(bl_rna.description)
else:
return ""
# Deselect all nodes in the tree.
@staticmethod
def deselect_nodes(context):
space = context.space_data
tree = space.edit_tree
for n in tree.nodes:
n.select = False
def create_node(self, context, node_type):
space = context.space_data
tree = space.edit_tree
try:
node = tree.nodes.new(type=node_type)
except RuntimeError as ex:
self.report({'ERROR'}, str(ex))
return None
node.select = True
tree.nodes.active = node
node.location = space.cursor_location
return node
def apply_node_settings(self, node):
for setting in self.settings:
# XXX catch exceptions here?
value = eval(setting.value)
node_data = node
node_attr_name = setting.name
# Support path to nested data.
if '.' in node_attr_name:
node_data_path, node_attr_name = node_attr_name.rsplit(".", 1)
node_data = node.path_resolve(node_data_path)
try:
setattr(node_data, node_attr_name, value)
except AttributeError as ex:
self.report(
{'ERROR_INVALID_INPUT'},
rpt_("Node has no attribute {:s}").format(setting.name))
print(str(ex))
# Continue despite invalid attribute
return node
# Base class for node "Add" operators.
class NodeAddOperator(NodeOperator):
use_transform: BoolProperty(
name="Use Transform",
description="Start transform operator after inserting the node",
default=False,
)
@staticmethod
def store_mouse_cursor(context, event):
space = context.space_data
tree = space.edit_tree
# convert mouse position to the View2D for later node placement
if context.region.type == 'WINDOW':
area = context.area
horizontal_pad = int(area.width / 10)
vertical_pad = int(area.height / 10)
inspace_x = min(max(horizontal_pad, event.mouse_region_x), area.width - horizontal_pad)
inspace_y = min(max(vertical_pad, event.mouse_region_y), area.height - vertical_pad)
# convert mouse position to the View2D for later node placement
space.cursor_location_from_region(inspace_x, inspace_y)
else:
space.cursor_location = tree.view_center
@classmethod
def poll(cls, context):
space = context.space_data
# needs active node editor and a tree to add nodes to
return (space and (space.type == 'NODE_EDITOR') and
space.edit_tree and space.edit_tree.is_editable)
# Default invoke stores the mouse position to place the node correctly
# and optionally invokes the transform operator
def invoke(self, context, event):
self.store_mouse_cursor(context, event)
result = self.execute(context)
if self.use_transform and ('FINISHED' in result):
# removes the node again if transform is canceled
bpy.ops.node.translate_attach_remove_on_cancel('INVOKE_DEFAULT')
return result
class NodeSwapOperator(NodeOperator):
properties_to_pass = (
'color',
'hide',
'label',
'mute',
'parent',
'show_options',
'show_preview',
'show_texture',
'use_alpha',
'use_clamp',
'use_custom_color',
"operation",
"domain",
"data_type",
)
@classmethod
def poll(cls, context):
if (context.area is None) or (context.area.type != "NODE_EDITOR"):
return False
if len(context.selected_nodes) <= 0:
cls.poll_message_set("No nodes selected.")
return False
return True
def transfer_node_properties(self, old_node, new_node):
for attr in self.properties_to_pass:
if (attr in self.settings):
return
if hasattr(old_node, attr) and hasattr(new_node, attr):
try:
setattr(new_node, attr, getattr(old_node, attr))
except (TypeError, ValueError):
pass
def transfer_input_values(self, old_node, new_node):
if (old_node.bl_idname in math_nodes) and (new_node.bl_idname in math_nodes):
for source_input, target_input in zip(old_node.inputs, new_node.inputs):
new_value = cast_value(source=source_input, target=target_input)
if new_value is not None:
target_input.default_value = new_value
else:
for input in old_node.inputs:
try:
new_socket = new_node.inputs[input.name]
new_value = cast_value(source=input, target=new_socket)
settings_name = "inputs[\"{:s}\"].default_value".format(input.name)
already_defined = (settings_name in self.settings)
if (new_value is not None) and not already_defined:
new_socket.default_value = new_value
except (AttributeError, KeyError, TypeError):
pass
@staticmethod
def transfer_links(tree, old_node, new_node, is_input):
both_math_nodes = (old_node.bl_idname in math_nodes) and (new_node.bl_idname in math_nodes)
if is_input:
if both_math_nodes:
for i, input in enumerate(old_node.inputs):
for link in input.links[:]:
try:
new_socket = new_node.inputs[i]
if new_socket.hide or not new_socket.enabled:
continue
tree.links.new(link.from_socket, new_socket)
except IndexError:
pass
else:
for input in old_node.inputs:
links = sorted(input.links, key=lambda link: link.multi_input_sort_id)
for link in links:
try:
new_socket = new_node.inputs[input.name]
if new_socket.hide or not new_socket.enabled:
continue
tree.links.new(link.from_socket, new_socket)
except KeyError:
pass
else:
if both_math_nodes:
for i, output in enumerate(old_node.outputs):
for link in output.links[:]:
try:
new_socket = new_node.outputs[i]
if new_socket.hide or not new_socket.enabled:
continue
new_link = tree.links.new(new_socket, link.to_socket)
except IndexError:
pass
else:
for output in old_node.outputs:
for link in output.links[:]:
try:
new_socket = new_node.outputs[output.name]
if new_socket.hide or not new_socket.enabled:
continue
new_link = tree.links.new(new_socket, link.to_socket)
try:
if link.to_socket.is_multi_input:
new_link.swap_multi_input_sort_id(link)
except AttributeError:
pass
except KeyError:
pass
@staticmethod
def get_switch_items(node):
switch_type = node.bl_idname
if switch_type == "GeometryNodeMenuSwitch":
return node.enum_definition.enum_items
if switch_type == "GeometryNodeIndexSwitch":
return node.index_switch_items
return None
def transfer_switch_data(self, old_node, new_node):
old_switch_items = self.get_switch_items(old_node)
new_switch_items = self.get_switch_items(new_node)
new_switch_items.clear()
if new_node.bl_idname == "GeometryNodeMenuSwitch":
for i, old_item in enumerate(old_switch_items[:]):
# Change the menu item names to numerical indices
# This makes it so that later functions that match by socket name work on the switches
if hasattr(old_item, "name"):
old_item.name = str(i)
new_switch_items.new(str(i))
if (old_switch_value := old_node.inputs[0].default_value) != '':
new_node.inputs[0].default_value = str(old_switch_value)
elif new_node.bl_idname == "GeometryNodeIndexSwitch":
for i, old_item in enumerate(old_switch_items[:]):
# Change the menu item names to numerical indices
# This makes it so that later functions that match by socket name work on the switches
if hasattr(old_item, "name"):
old_item.name = str(i)
new_switch_items.new()
if (old_switch_value := old_node.inputs[0].default_value) != '':
new_node.inputs[0].default_value = int(old_switch_value)
# Simple basic operator for adding a node.
class NODE_OT_add_node(NodeAddOperator, Operator):
"""Add a node to the active tree"""
bl_idname = "node.add_node"
bl_label = "Add Node"
bl_options = {'REGISTER', 'UNDO'}
type: StringProperty(
name="Node Type",
description="Node type",
)
visible_output: StringProperty(
name="Output Name",
description="If provided, all outputs that are named differently will be hidden",
options={'SKIP_SAVE'},
)
# Default execute simply adds a node.
def execute(self, context):
if self.properties.is_property_set("type"):
self.deselect_nodes(context)
if node := self.create_node(context, self.type):
self.apply_node_settings(node)
if self.visible_output:
for socket in node.outputs:
if socket.name != self.visible_output:
socket.hide = True
return {'FINISHED'}
else:
return {'CANCELLED'}
class NODE_OT_swap_node(NodeSwapOperator, Operator):
"""Replace the selected nodes with the specified type"""
bl_idname = "node.swap_node"
bl_label = "Swap Node"
bl_options = {"REGISTER", "UNDO"}
type: StringProperty(
name="Node Type",
description="Node type",
)
visible_output: StringProperty(
name="Output Name",
description="If provided, all outputs that are named differently will be hidden",
options={'SKIP_SAVE'},
)
@staticmethod
def get_zone_pair(tree, node):
# Get paired output node
if hasattr(node, "paired_output"):
return node, node.paired_output
# Get paired input node
for input_node in tree.nodes:
if hasattr(input_node, "paired_output"):
if input_node.paired_output == node:
return input_node, node
return None
def execute(self, context):
tree = context.space_data.edit_tree
for old_node in context.selected_nodes[:]:
if tree.nodes.get(old_node.name) is None:
continue
if old_node.bl_idname == self.type:
self.apply_node_settings(old_node)
continue
new_node = self.create_node(context, self.type)
self.apply_node_settings(new_node)
self.transfer_node_properties(old_node, new_node)
if self.visible_output:
for socket in new_node.outputs:
if socket.name != self.visible_output:
socket.hide = True
with temporary_unframe((old_node,)):
new_node.location = old_node.location
new_node.select = True
zone_pair = self.get_zone_pair(tree, old_node)
if zone_pair is not None:
input_node, output_node = zone_pair
if input_node.select and output_node.select:
with temporary_unframe((input_node, output_node)):
new_node.location = (input_node.location + output_node.location) / 2
new_node.select = True
self.transfer_input_values(input_node, new_node)
self.transfer_links(tree, input_node, new_node, is_input=True)
self.transfer_links(tree, output_node, new_node, is_input=False)
for node in zone_pair:
tree.nodes.remove(node)
else:
if (old_node.bl_idname in switch_nodes) and (new_node.bl_idname in switch_nodes):
self.transfer_switch_data(old_node, new_node)
self.transfer_input_values(old_node, new_node)
self.transfer_links(tree, old_node, new_node, is_input=True)
self.transfer_links(tree, old_node, new_node, is_input=False)
tree.nodes.remove(old_node)
return {'FINISHED'}
class NODE_OT_add_empty_group(NodeAddOperator, bpy.types.Operator):
bl_idname = "node.add_empty_group"
bl_label = "Add Empty Group"
bl_description = "Add a group node with an empty group"
bl_options = {'REGISTER', 'UNDO'}
# Override inherited method from NodeOperator
# Return None so that bl_description is used
@classmethod
def description(cls, _context, properties):
...
def execute(self, context):
from nodeitems_builtins import node_tree_group_type
tree = context.space_data.edit_tree
group = self.create_empty_group(tree.bl_idname)
self.deselect_nodes(context)
node = self.create_node(context, node_tree_group_type[tree.bl_idname])
self.apply_node_settings(node)
node.node_tree = group
return {"FINISHED"}
@staticmethod
def create_empty_group(idname):
group = bpy.data.node_groups.new(name="NodeGroup", type=idname)
input_node = group.nodes.new('NodeGroupInput')
input_node.select = False
input_node.location.x = -200 - input_node.width
output_node = group.nodes.new('NodeGroupOutput')
output_node.is_active_output = True
output_node.select = False
output_node.location.x = 200
return group
class NODE_OT_swap_empty_group(NodeSwapOperator, bpy.types.Operator):
bl_idname = "node.swap_empty_group"
bl_label = "Swap Empty Group"
bl_description = "Replace active node with an empty group"
bl_options = {'REGISTER', 'UNDO'}
# Override inherited method from NodeOperator
# Return None so that bl_description is used
@classmethod
def description(cls, _context, properties):
...
def execute(self, context):
from nodeitems_builtins import node_tree_group_type
tree = context.space_data.edit_tree
group = self.create_empty_group(tree.bl_idname)
bpy.ops.node.swap_node('INVOKE_DEFAULT', type=node_tree_group_type[tree.bl_idname])
for node in context.selected_nodes:
node.node_tree = group
return {"FINISHED"}
@staticmethod
def create_empty_group(idname):
group = bpy.data.node_groups.new(name="NodeGroup", type=idname)
input_node = group.nodes.new('NodeGroupInput')
input_node.select = False
input_node.location.x = -200 - input_node.width
output_node = group.nodes.new('NodeGroupOutput')
output_node.is_active_output = True
output_node.select = False
output_node.location.x = 200
return group
class ZoneOperator:
offset: FloatVectorProperty(
name="Offset",
description="Offset of nodes from the cursor when added",
size=2,
default=(150, 0),
)
_zone_tooltips = {
"GeometryNodeSimulationInput": (
"Simulate the execution of nodes across a time span"
),
"GeometryNodeRepeatInput": (
"Execute nodes with a dynamic number of repetitions"
),
"GeometryNodeForeachGeometryElementInput": (
"Perform operations separately for each geometry element (e.g. vertices, edges, etc.)"
),
"NodeClosureInput": (
"Wrap nodes inside a closure that can be executed at a different part of the nodetree"
),
}
@classmethod
def description(cls, _context, properties):
input_node_type = getattr(properties, "input_node_type", None)
# For Add Zone operators, use class variable instead of operator property
if input_node_type is None:
input_node_type = cls.input_node_type
return cls._zone_tooltips.get(input_node_type, None)
class NodeAddZoneOperator(ZoneOperator, NodeAddOperator):
add_default_geometry_link = True
def execute(self, context):
space = context.space_data
tree = space.edit_tree
self.deselect_nodes(context)
input_node = self.create_node(context, self.input_node_type)
output_node = self.create_node(context, self.output_node_type)
self.apply_node_settings(input_node)
self.apply_node_settings(output_node)
if input_node is None or output_node is None:
return {'CANCELLED'}
# Simulation input must be paired with the output.
input_node.pair_with_output(output_node)
input_node.location -= Vector(self.offset)
output_node.location += Vector(self.offset)
if tree.type == "GEOMETRY" and self.add_default_geometry_link:
# Connect geometry sockets by default if available.
# Get the sockets by their types, because the name is not guaranteed due to i18n.
from_socket = next(s for s in input_node.outputs if s.type == 'GEOMETRY')
to_socket = next(s for s in output_node.inputs if s.type == 'GEOMETRY')
tree.links.new(to_socket, from_socket)
return {'FINISHED'}
class NODE_OT_add_zone(NodeAddZoneOperator, Operator):
bl_idname = "node.add_zone"
bl_label = "Add Zone"
bl_options = {'REGISTER', 'UNDO'}
input_node_type: StringProperty(
name="Input Node",
description="Specifies the input node used the created zone",
)
output_node_type: StringProperty(
name="Output Node",
description="Specifies the output node used the created zone",
)
add_default_geometry_link: BoolProperty(
name="Add Geometry Link",
description="When enabled, create a link between geometry sockets in this zone",
default=False,
)
class NODE_OT_swap_zone(ZoneOperator, NodeSwapOperator, Operator):
bl_idname = "node.swap_zone"
bl_label = "Swap Zone"
bl_options = {"REGISTER", "UNDO"}
input_node_type: StringProperty(
name="Input Node",
description="Specifies the input node used the created zone",
)
output_node_type: StringProperty(
name="Output Node",
description="Specifies the output node used the created zone",
)
add_default_geometry_link: BoolProperty(
name="Add Geometry Link",
description="When enabled, create a link between geometry sockets in this zone",
default=False,
)
@staticmethod
def get_zone_pair(tree, node):
# Get paired output node
if hasattr(node, "paired_output"):
return node, node.paired_output
# Get paired input node
for input_node in tree.nodes:
if hasattr(input_node, "paired_output"):
if input_node.paired_output == node:
return input_node, node
return None
def execute(self, context):
tree = context.space_data.edit_tree
for old_node in context.selected_nodes[:]:
if tree.nodes.get(old_node.name) is None:
continue
zone_pair = self.get_zone_pair(tree, old_node)
if (old_node.bl_idname in {self.input_node_type, self.output_node_type}):
if zone_pair is not None:
old_input_node, old_output_node = zone_pair
self.apply_node_settings(old_input_node)
self.apply_node_settings(old_output_node)
else:
self.apply_node_settings(old_node)
continue
input_node = self.create_node(context, self.input_node_type)
output_node = self.create_node(context, self.output_node_type)
self.apply_node_settings(input_node)
self.apply_node_settings(output_node)
if input_node is None or output_node is None:
return {'CANCELLED'}
# Simulation input must be paired with the output.
input_node.pair_with_output(output_node)
if zone_pair is not None:
old_input_node, old_output_node = zone_pair
with temporary_unframe((old_input_node, old_output_node)):
input_node.location = old_input_node.location
output_node.location = old_output_node.location
self.transfer_node_properties(old_input_node, input_node)
self.transfer_node_properties(old_output_node, output_node)
self.transfer_input_values(old_input_node, input_node)
self.transfer_input_values(old_output_node, output_node)
self.transfer_links(tree, old_input_node, input_node, is_input=True)
self.transfer_links(tree, old_input_node, input_node, is_input=False)
self.transfer_links(tree, old_output_node, output_node, is_input=True)
self.transfer_links(tree, old_output_node, output_node, is_input=False)
for node in zone_pair:
tree.nodes.remove(node)
else:
with temporary_unframe((old_node,)):
input_node.location = old_node.location
output_node.location = old_node.location
input_node.location -= Vector(self.offset)
output_node.location += Vector(self.offset)
self.transfer_node_properties(old_node, input_node)
self.transfer_node_properties(old_node, output_node)
self.transfer_input_values(old_node, input_node)
self.transfer_links(tree, old_node, input_node, is_input=True)
self.transfer_links(tree, old_node, output_node, is_input=False)
tree.nodes.remove(old_node)
if tree.type == "GEOMETRY" and self.add_default_geometry_link:
# Connect geometry sockets by default if available.
# Get the sockets by their types, because the name is not guaranteed due to i18n.
from_socket = next(s for s in input_node.outputs if s.type == 'GEOMETRY')
to_socket = next(s for s in output_node.inputs if s.type == 'GEOMETRY')
if not (from_socket.is_linked or to_socket.is_linked):
tree.links.new(to_socket, from_socket)
return {'FINISHED'}
class NODE_OT_add_simulation_zone(NodeAddZoneOperator, Operator):
"""Add simulation zone input and output nodes to the active tree"""
bl_idname = "node.add_simulation_zone"
bl_label = "Add Simulation Zone"
bl_options = {'REGISTER', 'UNDO'}
input_node_type = "GeometryNodeSimulationInput"
output_node_type = "GeometryNodeSimulationOutput"
class NODE_OT_add_repeat_zone(NodeAddZoneOperator, Operator):
"""Add a repeat zone that allows executing nodes a dynamic number of times"""
bl_idname = "node.add_repeat_zone"
bl_label = "Add Repeat Zone"
bl_options = {'REGISTER', 'UNDO'}
input_node_type = "GeometryNodeRepeatInput"
output_node_type = "GeometryNodeRepeatOutput"
class NODE_OT_add_foreach_geometry_element_zone(NodeAddZoneOperator, Operator):
"""Add a For Each Geometry Element zone that allows executing nodes e.g. for each vertex separately"""
bl_idname = "node.add_foreach_geometry_element_zone"
bl_label = "Add For Each Geometry Element Zone"
bl_options = {'REGISTER', 'UNDO'}
input_node_type = "GeometryNodeForeachGeometryElementInput"
output_node_type = "GeometryNodeForeachGeometryElementOutput"
add_default_geometry_link = False
class NODE_OT_add_closure_zone(NodeAddZoneOperator, Operator):
"""Add a Closure zone"""
bl_idname = "node.add_closure_zone"
bl_label = "Add Closure Zone"
bl_options = {'REGISTER', 'UNDO'}
input_node_type = "NodeClosureInput"
output_node_type = "NodeClosureOutput"
add_default_geometry_link = False
class NODE_OT_collapse_hide_unused_toggle(Operator):
"""Toggle collapsed nodes and hide unused sockets"""
bl_idname = "node.collapse_hide_unused_toggle"
bl_label = "Collapse and Hide Unused Sockets"
bl_options = {'REGISTER', 'UNDO'}
@classmethod
def poll(cls, context):
space = context.space_data
# needs active node editor and a tree
return (space and (space.type == 'NODE_EDITOR') and
(space.edit_tree and space.edit_tree.is_editable))
def execute(self, context):
space = context.space_data
tree = space.edit_tree
for node in tree.nodes:
if node.select:
hide = (not node.hide)
node.hide = hide
# Note: connected sockets are ignored internally
for socket in node.inputs:
socket.hide = hide
for socket in node.outputs:
socket.hide = hide
return {'FINISHED'}
class NODE_OT_tree_path_parent(Operator):
"""Go to parent node tree"""
bl_idname = "node.tree_path_parent"
bl_label = "Parent Node Tree"
bl_options = {'REGISTER', 'UNDO'}
parent_tree_index: IntProperty(
name="Parent Index",
description="Parent index in context path",
default=0,
)
@classmethod
def poll(cls, context):
space = context.space_data
# needs active node editor and a tree
return (space and (space.type == 'NODE_EDITOR') and len(space.path) > 1)
def execute(self, context):
space = context.space_data
parent_number_to_pop = len(space.path) - 1 - self.parent_tree_index
for _ in range(parent_number_to_pop):
space.path.pop()
return {'FINISHED'}
class NodeInterfaceOperator():
@classmethod
def poll(cls, context):
space = context.space_data
if not space or space.type != 'NODE_EDITOR' or not space.edit_tree:
return False
if space.edit_tree.is_embedded_data:
return False
return True
class NODE_OT_interface_item_new(NodeInterfaceOperator, Operator):
"""Add a new item to the interface"""
bl_idname = "node.interface_item_new"
bl_label = "New Item"
bl_options = {'REGISTER', 'UNDO'}
def get_items(_self, context):
items = [
('INPUT', "Input", ""),
('OUTPUT', "Output", ""),
('PANEL', "Panel", ""),
]
if context is None:
return items
snode = context.space_data
tree = snode.edit_tree
interface = tree.interface
active_item = interface.active
# Panels have the extra option to add a toggle.
if active_item and active_item.item_type == 'PANEL':
items.append(('PANEL_TOGGLE', "Panel Toggle", ""))
return items
item_type: EnumProperty(
name="Item Type",
description="Type of the item to create",
items=get_items,
default=0,
)
# Returns a valid socket type for the given tree or None.
@staticmethod
def find_valid_socket_type(tree):
socket_type = 'NodeSocketFloat'
# Socket type validation function is only available for custom
# node trees. Assume that 'NodeSocketFloat' is valid for
# built-in node tree types.
if not hasattr(tree, "valid_socket_type") or tree.valid_socket_type(socket_type):
return socket_type
# Custom nodes may not support float sockets, search all
# registered socket subclasses.
types_to_check = [bpy.types.NodeSocket]
while types_to_check:
t = types_to_check.pop()
idname = getattr(t, "bl_idname", "")
if tree.valid_socket_type(idname):
return idname
# Test all subclasses
types_to_check.extend(t.__subclasses__())
def execute(self, context):
snode = context.space_data
tree = snode.edit_tree
interface = tree.interface
# Remember active item and position to determine target position.
active_item = interface.active
active_pos = active_item.position if active_item else -1
if self.item_type == 'INPUT':
item = interface.new_socket("Socket", socket_type=self.find_valid_socket_type(tree), in_out='INPUT')
elif self.item_type == 'OUTPUT':
item = interface.new_socket("Socket", socket_type=self.find_valid_socket_type(tree), in_out='OUTPUT')
elif self.item_type == 'PANEL':
item = interface.new_panel("Panel")
elif self.item_type == 'PANEL_TOGGLE':
active_panel = active_item
if len(active_panel.interface_items) > 0:
first_item = active_panel.interface_items[0]
if type(first_item) is bpy.types.NodeTreeInterfaceSocketBool and first_item.is_panel_toggle:
self.report({'INFO'}, "Panel already has a toggle")
return {'CANCELLED'}
item = interface.new_socket(active_panel.name, socket_type='NodeSocketBool', in_out='INPUT')
item.is_panel_toggle = True
interface.move_to_parent(item, active_panel, 0)
# Return in this case because we don't want to move the item.
return {'FINISHED'}
else:
return {'CANCELLED'}
if active_item:
# Insert into active panel if possible, otherwise insert after active item.
if active_item.item_type == 'PANEL' and item.item_type != 'PANEL':
interface.move_to_parent(item, active_item, len(active_item.interface_items))
else:
interface.move_to_parent(item, active_item.parent, active_pos + 1)
interface.active = item
return {'FINISHED'}
class NODE_OT_interface_item_duplicate(NodeInterfaceOperator, Operator):
"""Add a copy of the active item to the interface"""
bl_idname = "node.interface_item_duplicate"
bl_label = "Duplicate Item"
bl_options = {'REGISTER', 'UNDO'}
@classmethod
def poll(cls, context):
if not super().poll(context):
return False
snode = context.space_data
tree = snode.edit_tree
interface = tree.interface
return interface.active is not None
def execute(self, context):
snode = context.space_data
tree = snode.edit_tree
interface = tree.interface
item = interface.active
if item:
item_copy = interface.copy(item)
interface.active = item_copy
return {'FINISHED'}
class NODE_OT_interface_item_remove(NodeInterfaceOperator, Operator):
"""Remove active item from the interface"""
bl_idname = "node.interface_item_remove"
bl_label = "Remove Item"
bl_options = {'REGISTER', 'UNDO'}
def execute(self, context):
snode = context.space_data
tree = snode.edit_tree
interface = tree.interface
item = interface.active
if item:
if item.item_type == 'PANEL':
children = item.interface_items
if len(children) > 0:
first_child = children[0]
if isinstance(first_child, bpy.types.NodeTreeInterfaceSocket) and first_child.is_panel_toggle:
interface.remove(first_child)
interface.remove(item)
interface.active_index = min(interface.active_index, len(interface.items_tree) - 1)
# If the active selection lands on internal toggle socket, move selection to parent instead.
new_active = interface.active
if isinstance(new_active, bpy.types.NodeTreeInterfaceSocket) and new_active.is_panel_toggle:
interface.active_index = new_active.parent.index
return {'FINISHED'}
class NODE_OT_interface_item_make_panel_toggle(NodeInterfaceOperator, Operator):
"""Make the active boolean socket a toggle for its parent panel"""
bl_idname = "node.interface_item_make_panel_toggle"
bl_label = "Make Panel Toggle"
bl_options = {'REGISTER', 'UNDO'}
@classmethod
def poll(cls, context):
if not super().poll(context):
return False
snode = context.space_data
tree = snode.edit_tree
interface = tree.interface
active_item = interface.active
if not active_item:
return False
if type(active_item) is not bpy.types.NodeTreeInterfaceSocketBool or active_item.in_out != 'INPUT':
cls.poll_message_set("Only boolean input sockets are supported")
return False
parent_panel = active_item.parent
if parent_panel.parent is None:
cls.poll_message_set("Socket must be in a panel")
return False
if len(parent_panel.interface_items) > 0:
first_item = parent_panel.interface_items[0]
if first_item.is_panel_toggle:
cls.poll_message_set("Panel already has a toggle")
return False
return True
def execute(self, context):
snode = context.space_data
tree = snode.edit_tree
interface = tree.interface
active_item = interface.active
parent_panel = active_item.parent
if not parent_panel:
return {'CANCELLED'}
if type(active_item) is not bpy.types.NodeTreeInterfaceSocketBool:
return {'CANCELLED'}
active_item.is_panel_toggle = True
# Use the same name as the panel in the UI for clarity.
active_item.name = parent_panel.name
# Move the socket to the first position.
interface.move_to_parent(active_item, parent_panel, 0)
# Make the panel active.
interface.active = parent_panel
return {'FINISHED'}
class NODE_OT_interface_item_unlink_panel_toggle(NodeInterfaceOperator, Operator):
"""Make the panel toggle a stand-alone socket"""
bl_idname = "node.interface_item_unlink_panel_toggle"
bl_label = "Unlink Panel Toggle"
bl_options = {'REGISTER', 'UNDO'}
@classmethod
def poll(cls, context):
if not super().poll(context):
return False
snode = context.space_data
tree = snode.edit_tree
interface = tree.interface
active_item = interface.active
if not active_item or active_item.item_type != 'PANEL':
return False
if len(active_item.interface_items) == 0:
return False
first_item = active_item.interface_items[0]
return first_item.is_panel_toggle
def execute(self, context):
snode = context.space_data
tree = snode.edit_tree
interface = tree.interface
active_item = interface.active
if not active_item or active_item.item_type != 'PANEL':
return {'CANCELLED'}
if len(active_item.interface_items) == 0:
return {'CANCELLED'}
first_item = active_item.interface_items[0]
if type(first_item) is not bpy.types.NodeTreeInterfaceSocketBool or not first_item.is_panel_toggle:
return {'CANCELLED'}
first_item.is_panel_toggle = False
first_item.name = active_item.name
# Make the socket active.
interface.active = first_item
return {'FINISHED'}
class NODE_OT_viewer_shortcut_set(Operator):
"""Create a viewer shortcut for the selected node by pressing ctrl+1,2,..9"""
bl_idname = "node.viewer_shortcut_set"
bl_label = "Fast Preview"
bl_options = {'REGISTER', 'UNDO'}
viewer_index: IntProperty(
name="Viewer Index",
description="Index corresponding to the shortcut, e.g. number key 1 corresponds to index 1 etc..")
def get_connected_viewer(self, node):
for out in node.outputs:
for link in out.links:
nv = link.to_node
if nv.type == 'VIEWER':
return nv
return None
@classmethod
def poll(cls, context):
del cls
space = context.space_data
return (
(space is not None) and
space.type == 'NODE_EDITOR' and
space.node_tree is not None and
space.tree_type in {'CompositorNodeTree', 'GeometryNodeTree'}
)
def execute(self, context):
selected_nodes = context.selected_nodes
if len(selected_nodes) == 0:
self.report({'ERROR'}, "Select a node to assign a shortcut")
return {'CANCELLED'}
fav_node = selected_nodes[0]
# Only viewer nodes can be set to favorites. However, the user can
# create a new favorite viewer by selecting any node and pressing ctrl+1.
if fav_node.type == 'VIEWER':
viewer_node = fav_node
else:
viewer_node = self.get_connected_viewer(fav_node)
if not viewer_node:
# Calling `link_viewer()` if a viewer node is connected
# will connect the next available socket to the viewer node.
# This behavior is not desired as we want to create a shortcut to the existing connected viewer node.
# Therefore `link_viewer()` is called only when no viewer node is connected.
bpy.ops.node.link_viewer()
viewer_node = self.get_connected_viewer(fav_node)
if not viewer_node:
self.report(
{'ERROR'},
"Unable to set shortcut, selected node is not a viewer node or does not support viewing",
)
return {'CANCELLED'}
with bpy.context.temp_override(node=viewer_node):
bpy.ops.node.activate_viewer()
viewer_node.ui_shortcut = self.viewer_index
self.report({'INFO'}, rpt_("Assigned shortcut {:d} to {:s}").format(self.viewer_index, viewer_node.name))
return {'FINISHED'}
class NODE_OT_viewer_shortcut_get(Operator):
"""Toggle a specific viewer node using 1,2,..,9 keys"""
bl_idname = "node.viewer_shortcut_get"
bl_label = "Fast Preview"
bl_options = {'REGISTER', 'UNDO'}
viewer_index: IntProperty(
name="Viewer Index",
description="Index corresponding to the shortcut, e.g. number key 1 corresponds to index 1 etc..")
@classmethod
def poll(cls, context):
del cls
space = context.space_data
return (
(space is not None) and
space.type == 'NODE_EDITOR' and
space.node_tree is not None and
space.tree_type in {'CompositorNodeTree', 'GeometryNodeTree'}
)
def execute(self, context):
nodes = context.space_data.edit_tree.nodes
# Get viewer node with existing shortcut.
viewer_node = None
for n in nodes:
if n.type == 'VIEWER' and n.ui_shortcut == self.viewer_index:
viewer_node = n
if not viewer_node:
self.report({'INFO'}, rpt_("Shortcut {:d} is not assigned to a Viewer node yet").format(self.viewer_index))
return {'CANCELLED'}
with bpy.context.temp_override(node=viewer_node):
bpy.ops.node.toggle_viewer()
return {'FINISHED'}
class NODE_FH_image_node(FileHandler):
bl_idname = "NODE_FH_image_node"
bl_label = "Image node"
bl_import_operator = "node.add_image"
bl_file_extensions = ";".join((*bpy.path.extensions_image, *bpy.path.extensions_movie))
@classmethod
def poll_drop(cls, context):
return (
(context.area is not None) and
(context.area.type == 'NODE_EDITOR') and
(context.region is not None) and
(context.region.type == 'WINDOW')
)
classes = (
NodeSetting,
NODE_FH_image_node,
NODE_OT_add_node,
NODE_OT_swap_node,
NODE_OT_add_empty_group,
NODE_OT_swap_empty_group,
NODE_OT_add_zone,
NODE_OT_swap_zone,
NODE_OT_add_simulation_zone,
NODE_OT_add_repeat_zone,
NODE_OT_add_foreach_geometry_element_zone,
NODE_OT_add_closure_zone,
NODE_OT_collapse_hide_unused_toggle,
NODE_OT_interface_item_new,
NODE_OT_interface_item_duplicate,
NODE_OT_interface_item_remove,
NODE_OT_interface_item_make_panel_toggle,
NODE_OT_interface_item_unlink_panel_toggle,
NODE_OT_tree_path_parent,
NODE_OT_viewer_shortcut_get,
NODE_OT_viewer_shortcut_set,
)