Add more test cases for function node and fix it to work well with it
This commit is contained in:
StableLlama
@@ -23,11 +23,14 @@ class MathFormula(ComfyNodeABC):
|
||||
|
||||
This node takes an arbitrary number of numerical inputs (single letters like `a`, `b`, `c`, etc.) and safely
|
||||
evaluates the formula with supported operations: +, -, *, /, //, %, **, parentheses, and mathematical
|
||||
functions. It treats one-letter variables as inputs and multi-letter words as functions.
|
||||
functions.
|
||||
|
||||
NOTE on Operator Precedence: This parser uses standard precedence rules. Unary minus binds tightly,
|
||||
so expressions like `-a ** 2` are interpreted as `(-a) ** 2`. To calculate `-(a ** 2)`,
|
||||
use parentheses: `-(a ** 2)`.
|
||||
|
||||
The identifiers `e` and `pi` are special. When used as a function call (`e()`, `pi()`), they return their
|
||||
respective mathematical constants. When used as a variable (`e`), they are treated like any other
|
||||
variable and expect a corresponding input.
|
||||
respective mathematical constants. When used as a variable (`e`), they expect a corresponding input.
|
||||
|
||||
Supported functions:
|
||||
- Basic: abs, floor, ceil, round, min(a,b), max(a,b)
|
||||
@@ -42,9 +45,8 @@ class MathFormula(ComfyNodeABC):
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": ContainsDynamicDict({
|
||||
"formula": (IO.STRING, {"default": "a + b"}),
|
||||
"a": (IO.NUMBER, {"default": 1.0, "_dynamic": "letter"}),
|
||||
"b": (IO.NUMBER, {"default": 2.0, "_dynamic": "letter"}),
|
||||
"formula": (IO.STRING, {"default": "-e() ** -2"}),
|
||||
"a": (IO.NUMBER, {"default": 3.0, "_dynamic": "letter"}),
|
||||
}),
|
||||
}
|
||||
|
||||
@@ -62,9 +64,15 @@ class MathFormula(ComfyNodeABC):
|
||||
SUPPORTED_FUNCTIONS = set(FUNC_ARITIES.keys())
|
||||
|
||||
OPERATOR_PROPS = {
|
||||
"+": (2, "LEFT", 2), "-": (2, "LEFT", 2), "*": (3, "LEFT", 2), "/": (3, "LEFT", 2),
|
||||
"//": (3, "LEFT", 2), "%": (3, "LEFT", 2), "**": (4, "RIGHT", 2),
|
||||
"_NEG": (5, "RIGHT", 1),
|
||||
# token: (precedence, associativity, arity)
|
||||
"+": (2, "LEFT", 2),
|
||||
"-": (2, "LEFT", 2),
|
||||
"*": (3, "LEFT", 2),
|
||||
"/": (3, "LEFT", 2),
|
||||
"//": (3, "LEFT", 2),
|
||||
"%": (3, "LEFT", 2),
|
||||
"**": (4, "RIGHT", 2),
|
||||
"_NEG": (5, "RIGHT", 1), # Unary has higher precedence than exponentiation
|
||||
}
|
||||
OPERATORS = set(OPERATOR_PROPS.keys())
|
||||
|
||||
@@ -98,14 +106,13 @@ class MathFormula(ComfyNodeABC):
|
||||
output_queue.append(float(token))
|
||||
elif token.isalnum() and not token.isdigit():
|
||||
is_function_call = (i + 1 < len(tokens) and tokens[i+1] == '(')
|
||||
|
||||
if is_function_call:
|
||||
if token not in self.SUPPORTED_FUNCTIONS:
|
||||
raise ValueError(f"Unknown function: '{token}'")
|
||||
op_stack.append(token) # Push function name as string
|
||||
else: # It's a variable
|
||||
op_stack.append(token)
|
||||
else:
|
||||
if len(token) == 1 and token.isalpha():
|
||||
output_queue.append(('VAR', token)) # Push variable as a tagged tuple
|
||||
output_queue.append(('VAR', token))
|
||||
else:
|
||||
if token in self.SUPPORTED_FUNCTIONS:
|
||||
raise ValueError(f"Function '{token}' must be called with parentheses, e.g., {token}()")
|
||||
@@ -130,14 +137,17 @@ class MathFormula(ComfyNodeABC):
|
||||
elif token in self.OPERATORS:
|
||||
is_unary = token == '-' and (prev_token is None or prev_token in binary_operators or prev_token in ['(', ','])
|
||||
op_to_push = "_NEG" if is_unary else token
|
||||
|
||||
props = self.OPERATOR_PROPS[op_to_push]
|
||||
prec, assoc = props[0], props[1]
|
||||
prec = props[0]
|
||||
|
||||
while (op_stack and op_stack[-1] != '(' and op_stack[-1] in self.OPERATORS and
|
||||
(self.OPERATOR_PROPS.get(op_stack[-1], (0,))[0] > prec or
|
||||
(self.OPERATOR_PROPS.get(op_stack[-1], (0,))[0] == prec and assoc == "LEFT"))):
|
||||
output_queue.append(op_stack.pop())
|
||||
while (op_stack and op_stack[-1] != '(' and op_stack[-1] in self.OPERATORS):
|
||||
stack_op_props = self.OPERATOR_PROPS[op_stack[-1]]
|
||||
stack_op_prec = stack_op_props[0]
|
||||
stack_op_assoc = stack_op_props[1]
|
||||
if (stack_op_prec > prec) or (stack_op_prec == prec and stack_op_assoc == "LEFT"):
|
||||
output_queue.append(op_stack.pop())
|
||||
else:
|
||||
break
|
||||
op_stack.append(op_to_push)
|
||||
else:
|
||||
raise ValueError(f"Unknown identifier or invalid expression: '{token}'")
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import pytest
|
||||
import math
|
||||
import re
|
||||
from math import sin, cos
|
||||
from src.basic_data_handling.math_formula_node import MathFormula
|
||||
|
||||
def test_basic_formula_evaluation():
|
||||
@@ -31,6 +31,15 @@ def test_function_calls():
|
||||
# Single-argument functions
|
||||
assert node.evaluate("sin(a)", a=math.pi/6)[0] == pytest.approx(0.5)
|
||||
assert node.evaluate("cos(a)", a=0)[0] == pytest.approx(1.0)
|
||||
|
||||
# Unary minus with function calls
|
||||
assert node.evaluate("sin(-a)", a=math.pi/6)[0] == pytest.approx(-0.5)
|
||||
assert node.evaluate("-sin(a)", a=math.pi/6)[0] == pytest.approx(-0.5)
|
||||
assert node.evaluate("-sin(-a)", a=math.pi/6)[0] == pytest.approx(0.5)
|
||||
|
||||
# Unary minus with nested function calls
|
||||
assert node.evaluate("sin(cos(-a))", a=math.pi)[0] == pytest.approx(sin(cos(-math.pi)))
|
||||
assert node.evaluate("-sin(cos(-a))", a=math.pi)[0] == pytest.approx(-sin(cos(-math.pi)))
|
||||
assert node.evaluate("tan(a)", a=math.pi/4)[0] == pytest.approx(1.0)
|
||||
assert node.evaluate("sqrt(a)", a=16)[0] == pytest.approx(4.0)
|
||||
assert node.evaluate("log(a)", a=math.e)[0] == pytest.approx(1.0)
|
||||
@@ -75,6 +84,51 @@ def test_variable_constant_distinction():
|
||||
assert node.evaluate("e * log(e())", e=2)[0] == pytest.approx(2.0)
|
||||
assert node.evaluate("a * e() + b * e", a=2, b=3, e=4)[0] == pytest.approx(2 * math.e + 3 * 4)
|
||||
|
||||
def test_unary_minus():
|
||||
"""Test that unary minus (negation) is handled correctly in all contexts."""
|
||||
node = MathFormula()
|
||||
|
||||
# Simple variable negation
|
||||
assert node.evaluate("-a", a=5)[0] == pytest.approx(-5.0)
|
||||
assert node.evaluate("-a", a=-3)[0] == pytest.approx(3.0)
|
||||
|
||||
# Constant negation
|
||||
assert node.evaluate("-pi()")[0] == pytest.approx(-math.pi)
|
||||
assert node.evaluate("-e()")[0] == pytest.approx(-math.e)
|
||||
|
||||
# Function result negation
|
||||
assert node.evaluate("-sin(a)", a=math.pi/2)[0] == pytest.approx(-1.0)
|
||||
assert node.evaluate("-sqrt(a)", a=4)[0] == pytest.approx(-2.0)
|
||||
|
||||
# Nested unary minus
|
||||
assert node.evaluate("--a", a=5)[0] == pytest.approx(5.0)
|
||||
assert node.evaluate("---a", a=5)[0] == pytest.approx(-5.0)
|
||||
|
||||
# Unary minus with binary operators
|
||||
assert node.evaluate("a + -b", a=5, b=3)[0] == pytest.approx(2.0)
|
||||
assert node.evaluate("a * -b", a=4, b=3)[0] == pytest.approx(-12.0)
|
||||
assert node.evaluate("a / -b", a=6, b=2)[0] == pytest.approx(-3.0)
|
||||
assert node.evaluate("a - -b", a=5, b=3)[0] == pytest.approx(8.0)
|
||||
|
||||
# Unary minus with higher precedence operators
|
||||
assert node.evaluate("-a * b", a=2, b=3)[0] == pytest.approx(-6.0)
|
||||
assert node.evaluate("-(a * b)", a=2, b=3)[0] == pytest.approx(-6.0)
|
||||
assert node.evaluate("-a ** 2", a=3)[0] == pytest.approx(9.0) # see documentation about this special case
|
||||
assert node.evaluate("(-a) ** 2", a=3)[0] == pytest.approx(9.0) # (-a)^2
|
||||
|
||||
# Unary minus with parentheses
|
||||
assert node.evaluate("-(a + b)", a=2, b=3)[0] == pytest.approx(-5.0)
|
||||
assert node.evaluate("-(a + b) * c", a=2, b=3, c=4)[0] == pytest.approx(-20.0)
|
||||
assert node.evaluate("a * -(b + c)", a=2, b=3, c=4)[0] == pytest.approx(-14.0)
|
||||
|
||||
# Unary minus at start of parenthesized expressions
|
||||
assert node.evaluate("(-a + b)", a=5, b=3)[0] == pytest.approx(-2.0)
|
||||
assert node.evaluate("(-a - b)", a=5, b=3)[0] == pytest.approx(-8.0)
|
||||
|
||||
# Unary minus with constants and functions in complex expressions
|
||||
assert node.evaluate("-pi() * sin(-a)", a=math.pi/2)[0] == pytest.approx(math.pi)
|
||||
assert node.evaluate("-e() ** -2")[0] == pytest.approx(1/(math.e**2)) # see documentation about this special case
|
||||
|
||||
def test_error_conditions():
|
||||
"""Test error conditions in formula evaluation."""
|
||||
node = MathFormula()
|
||||
@@ -86,6 +140,20 @@ def test_error_conditions():
|
||||
with pytest.raises(ValueError, match=r"Variable 'e' was not provided."):
|
||||
node.evaluate("2 * e")
|
||||
|
||||
# Test unary minus with errors
|
||||
with pytest.raises(ValueError):
|
||||
node.evaluate("-unknown_var")
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
node.evaluate("-unknown_func(x)", x=5)
|
||||
|
||||
# Unary minus with domain errors
|
||||
with pytest.raises(ValueError):
|
||||
node.evaluate("sqrt(-a)", a=4) # Negative under square root
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
node.evaluate("log(-a)", a=2) # Negative in logarithm
|
||||
|
||||
# Missing variables
|
||||
with pytest.raises(ValueError, match=r"Variable 'x' was not provided"):
|
||||
node.evaluate("x + y", y=5)
|
||||
@@ -128,6 +196,18 @@ def test_complex_expressions():
|
||||
expected = 1 + 2 * 3 - 8 / 4 + 2 ** 3
|
||||
assert result == pytest.approx(expected)
|
||||
|
||||
# Complex expressions with unary minus
|
||||
formula = "a + -b * c - -(d / e) + -f ** g"
|
||||
result = node.evaluate(formula, a=1, b=2, c=3, d=8, e=4, f=2, g=3)[0]
|
||||
expected = 1 + (-2) * 3 - (-(8 / 4)) + (-(2 ** 3))
|
||||
assert result == pytest.approx(expected)
|
||||
|
||||
# Expression with multiple unary minuses and parentheses
|
||||
formula = "-a * (b + -c) / (-d)"
|
||||
result = node.evaluate(formula, a=2, b=5, c=3, d=4)[0]
|
||||
expected = -2 * (5 + (-3)) / (-4)
|
||||
assert result == pytest.approx(expected)
|
||||
|
||||
def test_formula_with_e_variable():
|
||||
"""Test formulas with 'e' as a variable to ensure it doesn't clash with the constant."""
|
||||
node = MathFormula()
|
||||
@@ -160,6 +240,14 @@ def test_whitespace_handling():
|
||||
# Lots of spaces
|
||||
assert node.evaluate("a + b * c", a=1, b=2, c=3)[0] == pytest.approx(7.0)
|
||||
|
||||
# Whitespace around unary minus
|
||||
assert node.evaluate("- a", a=5)[0] == pytest.approx(-5.0)
|
||||
assert node.evaluate("- a", a=5)[0] == pytest.approx(-5.0)
|
||||
assert node.evaluate("a * - b", a=3, b=4)[0] == pytest.approx(-12.0)
|
||||
assert node.evaluate("a * - b", a=3, b=4)[0] == pytest.approx(-12.0)
|
||||
assert node.evaluate("- (a + b)", a=2, b=3)[0] == pytest.approx(-5.0)
|
||||
assert node.evaluate("- (a + b)", a=2, b=3)[0] == pytest.approx(-5.0)
|
||||
|
||||
# Mixed spacing around functions
|
||||
assert node.evaluate("sin( a )+cos(b)", a=math.pi/2, b=0)[0] == pytest.approx(2.0)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user