armory/blender/arm/logicnode/math/LN_vector_math.py
2021-07-03 19:45:05 +02:00

155 lines
6.8 KiB
Python

from arm.logicnode.arm_nodes import *
class VectorMathNode(ArmLogicTreeNode):
"""Mathematical operations on vectors."""
bl_idname = 'LNVectorMathNode'
bl_label = 'Vector Math'
arm_section = 'vector'
arm_version = 1
min_outputs = 2
max_outputs = 5
def get_bool(self):
return self.get('property1', False)
def set_bool(self, value):
self['property1'] = value
if value:
if (self.property0 == 'Length') or (self.property0 == 'Distance') or (self.property0 == 'Dot Product'):
self.outputs.remove(self.outputs.values()[-1]) # Distance/Length/Scalar
self.add_output('NodeSocketFloat', 'X') # Result X
self.add_output('NodeSocketFloat', 'Y') # Result Y
self.add_output('NodeSocketFloat', 'Z') # Result Z
if (self.property0 == 'Length'):
self.add_output('NodeSocketFloat', 'Length') # Length
if (self.property0 == 'Distance'):
self.add_output('NodeSocketFloat', 'Distance') # Distance
if (self.property0 == 'Dot Product'):
self.add_output('NodeSocketFloat', 'Scalar') # Scalar
else:
if ((self.property0 == 'Length') or (self.property0 == 'Distance') or (self.property0 == 'Dot Product')) and (len(self.outputs) > 1):
self.outputs.remove(self.outputs.values()[-1]) # Distance/Length/Scalar
# Remove X, Y, Z
for i in range(3):
if len(self.outputs) > 1:
self.outputs.remove(self.outputs.values()[-1])
else:
break
if (self.property0 == 'Length'):
self.add_output('NodeSocketFloat', 'Length') # Length
if (self.property0 == 'Distance'):
self.add_output('NodeSocketFloat', 'Distance') # Distance
if (self.property0 == 'Dot Product'):
self.add_output('NodeSocketFloat', 'Scalar') # Scalar
property1: HaxeBoolProperty('property1', name='Separator Out', default=False, set=set_bool, get=get_bool)
@staticmethod
def get_enum_id_value(obj, prop_name, value):
return obj.bl_rna.properties[prop_name].enum_items[value].identifier
@staticmethod
def get_count_in(operation_name):
return {
'Add': 0,
'Subtract': 0,
'Average': 0,
'Dot Product': 0,
'Cross Product': 0,
'Multiply': 0,
'MultiplyFloats': 0,
'Distance': 2,
'Reflect': 2,
'Normalize': 1,
'Length': 1
}.get(operation_name, 0)
def get_enum(self):
return self.get('property0', 0)
def set_enum(self, value):
# Checking the selection of another operation
select_current = self.get_enum_id_value(self, 'property0', value)
select_prev = self.property0
if select_prev != select_current:
if (select_prev == 'Distance') or (select_prev == 'Length') or (self.property0 == 'Dot Product'):
self.outputs.remove(self.outputs.values()[-1])
# Many arguments: Add, Subtract, Average, Dot Product, Cross Product, Multiply, MultiplyFloats
if (self.get_count_in(select_current) == 0):
if (select_current == "MultiplyFloats") or (select_prev == "MultiplyFloats"):
while (len(self.inputs) > 1):
self.inputs.remove(self.inputs.values()[-1])
if (select_current == "MultiplyFloats"):
self.add_input('NodeSocketFloat', 'Value ' + str(len(self.inputs)))
else:
while (len(self.inputs) < 2):
self.add_input('NodeSocketVector', 'Value ' + str(len(self.inputs)))
if (select_current == 'Dot Product'):
self.add_output('NodeSocketFloat', 'Scalar')
# 2 arguments: Distance, Reflect
if (self.get_count_in(select_current) == 2):
count = 2
if (select_prev == "MultiplyFloats"):
count = 1
while (len(self.inputs) > count):
self.inputs.remove(self.inputs.values()[-1])
while (len(self.inputs) < 2):
self.add_input('NodeSocketVector', 'Value ' + str(len(self.inputs)))
if (select_current == 'Distance'):
self.add_output('NodeSocketFloat', 'Distance')
# 1 argument: Normalize, Length
if (self.get_count_in(select_current) == 1):
while (len(self.inputs) > 1):
self.inputs.remove(self.inputs.values()[-1])
if (select_current == 'Length'):
self.add_output('NodeSocketFloat', 'Length')
self['property0'] = value
property0: HaxeEnumProperty(
'property0',
items = [('Add', 'Add', 'Add'),
('Dot Product', 'Dot Product', 'Dot Product'),
('Multiply', 'Multiply', 'Multiply'),
('MultiplyFloats', 'Multiply (Floats)', 'Multiply (Floats)'),
('Normalize', 'Normalize', 'Normalize'),
('Subtract', 'Subtract', 'Subtract'),
('Average', 'Average', 'Average'),
('Cross Product', 'Cross Product', 'Cross Product'),
('Length', 'Length', 'Length'),
('Distance', 'Distance', 'Distance'),
('Reflect', 'Reflect', 'Reflect')],
name='', default='Add', set=set_enum, get=get_enum)
def __init__(self):
array_nodes[str(id(self))] = self
def init(self, context):
super(VectorMathNode, self).init(context)
self.add_input('NodeSocketVector', 'Value 0', default_value=[0.0, 0.0, 0.0])
self.add_input('NodeSocketVector', 'Value 1', default_value=[0.0, 0.0, 0.0])
self.add_output('NodeSocketVector', 'Result')
def draw_buttons(self, context, layout):
layout.prop(self, 'property1') # Separator Out
layout.prop(self, 'property0') # Operation
# Buttons
if (self.get_count_in(self.property0) == 0):
row = layout.row(align=True)
column = row.column(align=True)
op = column.operator('arm.node_add_input', text='Add Value', icon='PLUS', emboss=True)
op.node_index = str(id(self))
op.name_format = 'Value {0}'
if (self.property0 == "MultiplyFloats"):
op.socket_type = 'NodeSocketFloat'
else:
op.socket_type = 'NodeSocketVector'
column = row.column(align=True)
op = column.operator('arm.node_remove_input', text='', icon='X', emboss=True)
op.node_index = str(id(self))
if len(self.inputs) == 2:
column.enabled = False
def draw_label(self) -> str:
return f'{self.bl_label}: {self.property0}'