diff --git a/py/mathExpression.py b/py/mathExpression.py new file mode 100644 index 0000000..505a39a --- /dev/null +++ b/py/mathExpression.py @@ -0,0 +1,95 @@ +import ast +import operator as op + +operators = { + ast.Add: op.add, + ast.Sub: op.sub, + ast.Mult: op.mul, + ast.Div: op.truediv, + ast.FloorDiv: op.floordiv, + ast.Pow: op.pow, + ast.BitXor: op.xor, + ast.USub: op.neg, + ast.Mod: op.mod +} + + +class MathExpression: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "expression": ("STRING", {"multiline": True, "dynamicPrompts": False}), + }, + "optional": { + "a": ("FLOAT", {"forceInput": True}), + "b": ("FLOAT", {"forceInput": True}), + "c": ("FLOAT", {"forceInput": True}), + }, + "hidden": {"extra_pnginfo": "EXTRA_PNGINFO", + "prompt": "PROMPT"}, + } + + RETURN_TYPES = ("INT", "FLOAT", "STRING", ) + FUNCTION = "evaluate" + CATEGORY = "util" + + def get_widget_value(self, extra_pnginfo, prompt, node_name, widget_name): + workflow = extra_pnginfo["workflow"] + node_id = None + for node in workflow["nodes"]: + name = node["type"] + if "properties" in node: + if "Node name for S&R" in node["properties"]: + name = node["properties"]["Node name for S&R"] + if name == node_name: + node_id = node["id"] + break + if "title" in node: + name = node["title"] + if name == node_name: + node_id = node["id"] + break + if node_id is not None: + values = prompt[str(node_id)] + if "inputs" in values: + if widget_name in values["inputs"]: + return values["inputs"][widget_name] + raise NameError(f"Widget not found: {node_name}.{widget_name}") + raise NameError(f"Node not found: {node_name}.{widget_name}") + + def evaluate(self, expression, extra_pnginfo, prompt, a=None, b=None, c=None): + node = ast.parse(expression, mode='eval').body + + def eval_expr(node): + if isinstance(node, ast.Num): + return node.n + elif isinstance(node, ast.BinOp): + return operators[type(node.op)](eval_expr(node.left), eval_expr(node.right)) + elif isinstance(node, ast.UnaryOp): + return operators[type(node.op)](eval_expr(node.operand)) + elif isinstance(node, ast.Attribute): + return self.get_widget_value(extra_pnginfo, prompt, node.value.id, node.attr) + elif isinstance(node, ast.Name): + if node.id == "a": + return a + if node.id == "b": + return b + if node.id == "c": + return c + raise NameError(f"Name not found: {node.id}") + else: + raise TypeError(node) + + r = eval_expr(node) + return (int(r), float(r), str(r)) + + +NODE_CLASS_MAPPINGS = { + "MathExpression|pysssss": MathExpression, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "MathExpression|pysssss": "Math Expression 🐍", +}