Borealis-Legacy/Nodes/math_operation_node.py

122 lines
4.3 KiB
Python

#!/usr/bin/env python3
"""
Math Operation Node:
- Inputs: Two input ports ("A" and "B").
- Output: One output port ("Result").
- Operation: A dropdown (combo menu) to select:
Add, Subtract, Multiply, Divide, Average.
- Displays the computed result in a read-only text box labeled "Result".
"""
from NodeGraphQt import BaseNode
class MathOperationNode(BaseNode):
__identifier__ = 'bunny-lab.io.math_node'
NODE_NAME = 'Math Operation'
def __init__(self):
super(MathOperationNode, self).__init__()
# ----------------------------------------------------------------------
# Initialization Section:
# - Create two input ports: A, B
# - Create one output port: Result
# - Add a combo box for operator selection
# - Add a text input for displaying the computed result
# ----------------------------------------------------------------------
self.add_input('A')
self.add_input('B')
self.add_output('Result')
# Operator combo box (Add, Subtract, Multiply, Divide, Average)
self.add_combo_menu('operator', 'Operator', items=[
'Add', 'Subtract', 'Multiply', 'Divide', 'Average'
])
# Text input for displaying the computed result.
# We'll make it read-only by accessing the underlying QLineEdit.
self.add_text_input('calc_result', 'Result', text='0')
result_widget = self.get_widget('calc_result') # This is a NodeLineEdit wrapper
if result_widget:
# Get the underlying QLineEdit
line_edit = result_widget.get_custom_widget()
# Make the QLineEdit read-only
line_edit.setReadOnly(True)
self.value = 0
self.set_name("Math Operation")
self.process_input()
def process_input(self, event=None):
"""
Compute Section:
- For each input port (A, B), if connected, grab the 'value' from
the upstream node; otherwise default to 0.0.
- Convert to float, apply the selected operation from the combo box,
update the "Result" text box, node title, and output port.
"""
# Gather input A
input_a = self.input(0)
if input_a and input_a.connected_ports():
a_raw = input_a.connected_ports()[0].node().get_property('value')
else:
a_raw = 0.0
# Gather input B
input_b = self.input(1)
if input_b and input_b.connected_ports():
b_raw = input_b.connected_ports()[0].node().get_property('value')
else:
b_raw = 0.0
# Convert raw inputs to floats (default 0.0 on failure).
try:
a_val = float(a_raw)
except (ValueError, TypeError):
a_val = 0.0
try:
b_val = float(b_raw)
except (ValueError, TypeError):
b_val = 0.0
# Retrieve the selected operator from the combo box.
operator = self.get_property('operator')
result = 0.0
if operator == 'Add':
result = a_val + b_val
elif operator == 'Subtract':
result = a_val - b_val
elif operator == 'Multiply':
result = a_val * b_val
elif operator == 'Divide':
result = a_val / b_val if b_val != 0 else 0.0
elif operator == 'Average':
result = (a_val + b_val) / 2.0
self.value = result
# Update the read-only text input and node title.
self.set_property('calc_result', str(result))
# Transmit the numeric result to any connected output nodes.
output_port = self.output(0)
if output_port and output_port.connected_ports():
for cp in output_port.connected_ports():
connected_node = cp.node()
if hasattr(connected_node, 'receive_data'):
connected_node.receive_data(result, source_port_name='Result')
def on_input_connected(self, input_port, output_port):
self.process_input()
def on_input_disconnected(self, input_port, output_port):
self.process_input()
def property_changed(self, property_name):
if property_name in ['operator']:
self.process_input()
def receive_data(self, data, source_port_name=None):
self.process_input()