diff --git a/.gitignore b/.gitignore index 0a2a367..1526efa 100644 --- a/.gitignore +++ b/.gitignore @@ -8,4 +8,6 @@ rgthree_config.json web/rgthree_config.js web/comfyui/rgthree_config.js userdata/ -userdata/** \ No newline at end of file +userdata/** +web/comfyui/testing/ +web/comfyui/tests/ \ No newline at end of file diff --git a/__build__.py b/__build__.py index afa7720..2410ab8 100644 --- a/__build__.py +++ b/__build__.py @@ -46,7 +46,7 @@ def log_step_info(msg:str, status='info'): step_infos.append({"msg": f' - {msg}', "type": status}) -def build(with_tests = False, fix = False): +def build(without_tests = True, fix = False): THIS_DIR = os.path.dirname(os.path.abspath(__file__)) DIR_SRC_WEB = os.path.abspath(f'{THIS_DIR}/src_web/') @@ -84,10 +84,10 @@ def build(with_tests = False, fix = False): checked = subprocess.run(["node", "./node_modules/typescript/bin/tsc"], check=True) log_step(status="Done") - if with_tests: + if not without_tests: log_step(msg='Removing directories (KEEPING TESTING)', status="Notice") else: - log_step(msg='Removing uneeded directories') + log_step(msg='Removing unneeded directories') test_path = os.path.join(DIR_WEB, 'comfyui', 'tests') if os.path.exists(test_path): rmtree(test_path) @@ -137,10 +137,10 @@ def build(with_tests = False, fix = False): if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("-t", "--with-tests", default=False, action="store_true") + parser.add_argument("-t", "--no-tests", default=False, action="store_true") parser.add_argument("-f", "--fix", default=False, action="store_true") args = parser.parse_args() start = time.time() - build(with_tests=args.with_tests, fix=args.fix) + build(without_tests=args.no_tests, fix=args.fix) print(f'Finished all in {round(time.time() - start, 3)}s') diff --git a/py/power_puter.py b/py/power_puter.py index 2c4b129..1c94e3d 100644 --- a/py/power_puter.py +++ b/py/power_puter.py @@ -17,6 +17,7 @@ from typing import Any, Callable import operator as op from .constants import get_category, get_name from .utils import FlexibleOptionalInputType, any_type, get_dict_value +from .log import log_node_warn @dataclasses.dataclass(frozen=True, kw_only=True) @@ -44,6 +45,7 @@ _FUNCTIONS = { Function(name="max", call=max, args=(2, None)), Function(name="random_int", call=random.randint, args=(2, 2)), Function(name="random_choice", call=random.choice, args=(2, None)), + Function(name="re", call=re.compile, args=(1, 1)), # Casts Function(name="int", call=int, args=(1, 1)), Function(name="float", call=float, args=(1, 1)), @@ -62,7 +64,6 @@ _FUNCTIONS = { # tell. _GLOBAL_FUNCTION_NAMES = ['node', 'nodes'] - # Special functions by class type (called from the Attrs.) _SPECIAL_FUNCTIONS = { "Power Lora Loader (rgthree)": { @@ -79,7 +80,6 @@ _SPECIAL_FUNCTIONS = { } } - _OPERATORS = { ast.Add: op.add, ast.Sub: op.sub, @@ -102,11 +102,13 @@ _OPERATORS = { _IS_CHANGED_GLOBAL = 0 +_NODE_NAME = get_name("Power Puter") + class RgthreePowerPuter: """A powerful node that can compute and evaluate expressions and output as various types.""" - NAME = get_name("Power Puter") + NAME = _NODE_NAME CATEGORY = get_category() @classmethod @@ -178,6 +180,9 @@ class _Puter: self._code = code self._workflow = workflow self._prompt = prompt + self._prompt_nodes = [] + if self._prompt: + self._prompt_nodes = [{'id': id} | {**node} for id, node in self._prompt.items()] def execute(self, code=str | None) -> Any: """Evaluates a the code block.""" @@ -191,26 +196,33 @@ class _Puter: break return last_value - def _get_nodes(self, node_id: int | str | None = None) -> dict[str, Any]: - """Get a dict of the nodes that match the node_id, or all the nodes in the prompt.""" + def _get_nodes(self, node_id: int | str | re.Pattern | None = None) -> list[Any]: + """Get a list of the nodes that match the node_id, or all the nodes in the prompt.""" + log_node_warn( + _NODE_NAME, + "There was a breaking change to `nodes()` in Power Puter. The return is now a list of nodes," + " rather than a dict. Update any Power Puter code." + ) + nodes = self._prompt_nodes.copy() if not node_id: - return {**self._prompt} - node_id = str(node_id) - nodes = [] - if re.match(r'\d+$', node_id): - nodes = {id: n for id, n in self._prompt.items() if node_id == id} - if not nodes: - nodes = { - id: n for id, n in self._prompt.items() if node_id == get_dict_value(n, '_meta.title', '') - } - return nodes + return nodes - def _get_node(self, node_id: int | str): + if isinstance(node_id, re.Pattern): + found = [n for n in nodes if re.search(node_id, get_dict_value(n, '_meta.title', ''))] + else: + node_id = str(node_id) + found = None + if re.match(r'\d+$', node_id): + found = [n for n in nodes if node_id == n['id']] + if not found: + found = [n for n in nodes if node_id == get_dict_value(n, '_meta.title', '')] + return found + + def _get_node(self, node_id: int | str | re.Pattern) -> Any | None: """Returns a prompt-node from the hidden prompt.""" - node_id = str(node_id) - nodes = [n for n in self._get_nodes(node_id).values()] + nodes = self._get_nodes(node_id) if nodes and len(nodes) > 1: - print('ERROR more than one node, returning first.') + log_node_warn(_NODE_NAME, f"More than one node found for '{node_id}'. Returning first.") return nodes[0] if nodes else None def _eval_statement(self, stmt: ast.stmt, ctx: dict | None = None): @@ -285,42 +297,56 @@ class _Puter: # Like: [v.lower() for v in lora_list] # Like: [v for v in l if v.startswith('B')] # Like: [v.lower() for v in l if v.startswith('B') or v.startswith('F')] + # --- + # Like: [l for n in nodes(re('Loras')).values() if (l := n.loras)] final_list = [] - for gen in stmt.generators: - gen_ctx = {**ctx} - if isinstance(gen.target, ast.Tuple): - gen_ctx[gen.target.elts[0].id] = None - gen_ctx[gen.target.elts[1].id] = None - elif isinstance(gen.target, ast.Name): + gen_ctx = {**ctx} + + generators = [*stmt.generators] + + def handle_gen(generators: list[ast.comprehension]): + gen = generators.pop(0) + if isinstance(gen.target, ast.Name): gen_ctx[gen.target.id] = None + elif isinstance(gen.target, ast.Tuple): # dict, like `for k, v in d.entries()` + for elt in gen.target.elts: + gen_ctx[elt.id] = None + else: + raise ValueError('Na') # A call, like my_dct.items(), or a named ctx list if isinstance(gen.iter, ast.Call): - iter = self._eval_statement(gen.iter.func, ctx=gen_ctx)() + gen_iters = self._eval_statement(gen.iter.func, ctx=gen_ctx)() elif isinstance(gen.iter, (ast.Name, ast.Attribute)): - iter = self._eval_statement(gen.iter, ctx=gen_ctx) + gen_iters = self._eval_statement(gen.iter, ctx=gen_ctx) - for v in iter: - # Unpack if we were a dict, otherwise it's just the item - if isinstance(v, (tuple, list)): - gen_ctx[gen.target.elts[0].id] = v[0] - gen_ctx[gen.target.elts[1].id] = v[1] + for gen_iter in gen_iters: + if_ctx = {**gen_ctx} + if isinstance(gen.target, ast.Tuple): # dict, like `for k, v in d.entries()` + for i, elt in enumerate(gen.target.elts): + if_ctx[elt.id] = gen_iter[i] else: - gen_ctx[gen.target.id] = v - + if_ctx[gen.target.id] = gen_iter good = True for ifcall in gen.ifs: - if not self._eval_statement(ifcall, ctx=gen_ctx): + if not self._eval_statement(ifcall, ctx=if_ctx): good = False break if not good: continue + gen_ctx.update(if_ctx) + if len(generators): + handle_gen(generators) + else: + final_list.append(self._eval_statement(stmt.elt, gen_ctx)) + generators.insert(0, gen) - final_list.append(self._eval_statement(stmt.elt, gen_ctx)) - return final_list + handle_gen(generators) + return final_list if isinstance(stmt, ast.Call): + call = None if isinstance(stmt.func, ast.Attribute): call = self._eval_statement(stmt.func, ctx=ctx) if isinstance(stmt.func, ast.Name): @@ -335,6 +361,8 @@ class _Puter: toErr = " or more" if fn.args[1] is None else f" to {fn.args[1]}" raise SyntaxError(f"Invalid function call: {fn.name} requires {fn.args[0]}{toErr} args") if not call: + # print(type(stmt)) + # print(type(stmt.func)) raise ValueError(f'No call for ast.Call {stmt}') args = [] for arg in stmt.args: @@ -363,7 +391,14 @@ class _Puter: # Assign a variable and add it to our ctx. if isinstance(stmt, ast.Assign): value = self._eval_statement(stmt.value, ctx=ctx) - self._ctx[stmt.targets[0].id] = value + ctx[stmt.targets[0].id] = value + return value + + # For assigning a var in a list comprehension. + # Like [name for node in node_list if (name := node.name)] + if isinstance(stmt, ast.NamedExpr): + value = self._eval_statement(stmt.value, ctx=ctx) + ctx[stmt.target.id] = value return value raise TypeError(stmt) diff --git a/pyproject.toml b/pyproject.toml index b335c5e..e058c88 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "rgthree-comfy" description = "Making ComfyUI more comfortable." -version = "1.0.2505182359" +version = "1.0.2505212110" license = { file = "LICENSE" } dependencies = [] diff --git a/src_web/comfyui/tests/power_puter.ts b/src_web/comfyui/tests/power_puter.ts new file mode 100644 index 0000000..6caa6cb --- /dev/null +++ b/src_web/comfyui/tests/power_puter.ts @@ -0,0 +1,56 @@ +import type {LGraphNode} from "@comfyorg/litegraph"; + +import {NodeTypesString} from "../constants"; +import {ComfyUITestEnvironment} from "../testing/comfyui_env"; +import {describe, should, beforeEach, expect, describeRun} from "../testing/runner.js"; + +const env = new ComfyUITestEnvironment(); + +function setPowerPuterValue(node: LGraphNode, value: string, outputType: string = "STRING") { + node.widgets![1]!.value = value; + node.widgets![0]!.value = outputType; +} + +describe("TestPowerPuter", async () => { + let powerPuter!: LGraphNode; + let displayAny!: LGraphNode; + + await beforeEach(async () => { + await env.clear(); + powerPuter = await env.addNode(NodeTypesString.POWER_PUTER); + displayAny = await env.addNode(NodeTypesString.DISPLAY_ANY); + powerPuter.connect(0, displayAny, 0); + await env.wait(); + }); + + await should("output constants and concatenation", async () => { + const checks: Array<[string, string] | [string, string, string]> = [ + ["1", "1"], + ['"abc"', "abc"], + ["1 + 2", "3"], + ['"abc" + "xyz"', "abcxyz"], + // INT + ["1", "1", "INT"], + ["1 + 2", "3", "INT"], + // FLOAT + ["1", "1.0", "FLOAT"], + ["1.3 + 2.8", "4.1", "FLOAT"], + // BOOL + ["1", "True", "BOOL"], + ["1 - 1", "False", "BOOL"], + ]; + for (const data of checks) { + setPowerPuterValue(powerPuter, data[0], data[2]); + await env.queuePrompt(); + expect(displayAny.widgets![0]!.value).toBe(data[0], data[1]); + } + }); + + await should("handle inputs", async () => { + // TODO + }); + + await should("handle complex inputs", async () => { + // TODO + }); +});