diff --git a/more_math/ScriptTextWindow.py b/more_math/ScriptTextWindow.py index 5be2781..9d52f06 100644 --- a/more_math/ScriptTextWindow.py +++ b/more_math/ScriptTextWindow.py @@ -22,4 +22,4 @@ class ScriptTextInput(io.ComfyNode): @classmethod def execute(cls, script): - return (remove_dead_code(parse_expr(script)),) + return (remove_dead_code(parse_expr.__wrapped__(script)),) diff --git a/more_math/helper_functions.py b/more_math/helper_functions.py index 621f3bc..aa1a96a 100644 --- a/more_math/helper_functions.py +++ b/more_math/helper_functions.py @@ -1,3 +1,5 @@ +import functools + from antlr4.error.ErrorListener import ErrorListener from antlr4 import InputStream, CommonTokenStream from antlr4.atn.PredictionMode import PredictionMode @@ -66,8 +68,10 @@ def as_tensor(value, shape): return torch.cat(value) +# check_lazy_status and execute both parse the same expression; trees are only read, so they can be shared. +@functools.lru_cache(maxsize=64) def parse_expr(expr: str): - """Parse a math expression and return the parse tree.""" + """Parse a math expression and return the parse tree. The tree is shared, use parse_expr.__wrapped__ for one you modify.""" input_stream = InputStream(expr) lexer = MathExprLexer(input_stream) stream = CommonTokenStream(lexer) diff --git a/tests/test_dead_code.py b/tests/test_dead_code.py index 6fd3d70..1080918 100644 --- a/tests/test_dead_code.py +++ b/tests/test_dead_code.py @@ -9,7 +9,7 @@ from more_math.helper_functions import parse_expr, checkLazyNew def run(script, variables=None, optimize=True, stack=None): - tree = parse_expr(script) + tree = parse_expr.__wrapped__(script) if optimize: remove_dead_code(tree) out = io.StringIO() @@ -109,5 +109,5 @@ def test_last_statement_of_a_block_stays(): def test_unused_inputs_are_no_longer_requested(): - tree = remove_dead_code(parse_expr("u = V1 * 2; V0")) + tree = remove_dead_code(parse_expr.__wrapped__("u = V1 * 2; V0")) assert checkLazyNew(tree, {"V0": None, "V1": None}, {}) == {"V0"}