diff --git a/py/power_puter.py b/py/power_puter.py index 669d40a..ca04135 100644 --- a/py/power_puter.py +++ b/py/power_puter.py @@ -510,11 +510,9 @@ class _Puter: else: # Slice could be a name or a constant; evaluate it attr = self._eval_statement(stmt.slice, ctx=ctx) - if not isinstance(attr, str): - raise ValueError(f'Disallowed access to "{attr}"') # Check if we're blocking access to this attribute/method on this item type. for typ, names in _BLOCKED_METHODS_OR_ATTRS.items(): - if isinstance(item, typ) and attr in names: + if isinstance(item, typ) and isinstance(attr, str) and attr in names: raise ValueError(f'Disallowed access to "{attr}" for type {typ}.') try: val = item[attr] diff --git a/pyproject.toml b/pyproject.toml index 0680d84..7a6ad66 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "rgthree-comfy" description = "Making ComfyUI more comfortable." -version = "1.0.2512060053" +version = "1.0.2512071717" license = { file = "LICENSE" } dependencies = [] diff --git a/src_web/comfyui/testing/runner.ts b/src_web/comfyui/testing/runner.ts index 0be2023..0e9a4db 100644 --- a/src_web/comfyui/testing/runner.ts +++ b/src_web/comfyui/testing/runner.ts @@ -86,6 +86,7 @@ function log(msg: string, styles: string) { class Expectation { private propertyLabel: string | null = ""; private expectedLabel: string | null = ""; + private verbLabel: string | null = "be"; private expectedFn!: (v: any) => boolean; private value: any; @@ -124,9 +125,17 @@ class Expectation { this.expectedLabel = "a number"; return this.toBeEval(); } + toContain(labelOrExpected: any, maybeExpected?: any) { + const expected = maybeExpected !== undefined ? maybeExpected : labelOrExpected; + this.propertyLabel = maybeExpected !== undefined ? labelOrExpected : null; + this.verbLabel = 'contain'; + this.expectedLabel = JSON.stringify(expected); + this.expectedFn = (v) => v.includes(expected); + return this.toBeEval(); + } toBeEval(strict = false) { let evaluation = this.expectedFn(this.value); - let msg = `Expected ${this.propertyLabel ? this.propertyLabel + " to be " : ""}${ + let msg = `Expected ${this.propertyLabel ? this.propertyLabel + ` to ${this.verbLabel} ` : ""}${ this.expectedLabel }`; msg += evaluation ? "." : `, but was ${JSON.stringify(this.value)}`; diff --git a/src_web/comfyui/testing/utils_test.ts b/src_web/comfyui/testing/utils_test.ts new file mode 100644 index 0000000..9c805f2 --- /dev/null +++ b/src_web/comfyui/testing/utils_test.ts @@ -0,0 +1,41 @@ +import type {LGraphNode} from "@comfyorg/frontend"; +import type {ComfyUITestEnvironment} from "./comfyui_env"; + +export const PNG_1x1 = + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVQIW2P4v5ThPwAG7wKklwQ/bwAAAABJRU5ErkJggg=="; +export const PNG_1x2 = + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAACCAYAAACZgbYnAAAAEElEQVQIW2NgYGD4D8QM/wEHAwH/OMSHKAAAAABJRU5ErkJggg=="; +export const PNG_2x1 = + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAIAAAABCAYAAAD0In+KAAAAD0lEQVQIW2NkYGD4D8QMAAUNAQFqjhCLAAAAAElFTkSuQmCC"; + +export async function pasteImageToLoadImageNode( + env: ComfyUITestEnvironment, + dataUrl?: string, + node?: LGraphNode, +) : Promise { + const dataArr = (dataUrl ?? PNG_1x1).split(","); + const mime = dataArr[0]!.match(/:(.*?);/)![1]; + const bstr = atob(dataArr[1]!); + let n = bstr.length; + const u8arr = new Uint8Array(n); + while (n--) { + u8arr[n] = bstr.charCodeAt(n); + } + const filename = `test_image_${+new Date()}.png`; + const file = new File([u8arr], filename, {type: mime}); + if (!node) { + node = await env.addNode("LoadImage"); + } + await (node as any).pasteFiles([file]); + let i = 0; + let good = false; + while (i++ < 10 || good) { + good = node.widgets![0]!.value === filename; + if (good) break; + await env.wait(100); + } + if (!good) { + throw new Error("Expected file not loaded."); + } + return node; +} diff --git a/src_web/comfyui/tests/image_or_latent_size_tests.ts b/src_web/comfyui/tests/image_or_latent_size_tests.ts index 3c35972..0cb932b 100644 --- a/src_web/comfyui/tests/image_or_latent_size_tests.ts +++ b/src_web/comfyui/tests/image_or_latent_size_tests.ts @@ -3,41 +3,10 @@ import type {LGraphNode} from "@comfyorg/frontend"; import {NodeTypesString} from "../constants"; import {ComfyUITestEnvironment} from "../testing/comfyui_env"; import {describe, should, beforeEach, expect, describeRun} from "../testing/runner.js"; +import {pasteImageToLoadImageNode, PNG_1x1, PNG_1x2, PNG_2x1} from "../testing/utils_test.js"; const env = new ComfyUITestEnvironment(); -const PNG_1x1 = - "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVQIW2P4v5ThPwAG7wKklwQ/bwAAAABJRU5ErkJggg=="; -const PNG_1x2 = - "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAACCAYAAACZgbYnAAAAEElEQVQIW2NgYGD4D8QM/wEHAwH/OMSHKAAAAABJRU5ErkJggg=="; -const PNG_2x1 = - "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAIAAAABCAYAAAD0In+KAAAAD0lEQVQIW2NkYGD4D8QMAAUNAQFqjhCLAAAAAElFTkSuQmCC"; - -async function pastImageToLoadImgeNode(dataUrl: string, node: LGraphNode) { - const dataArr = dataUrl.split(","); - const mime = dataArr[0]!.match(/:(.*?);/)![1]; - const bstr = atob(dataArr[1]!); - let n = bstr.length; - const u8arr = new Uint8Array(n); - while (n--) { - u8arr[n] = bstr.charCodeAt(n); - } - const filename = `test_image_${+new Date()}.png`; - const file = new File([u8arr], filename, {type: mime}); - await (node as any).pasteFiles([file]); - let i = 0; - let good = false; - while (i++ < 10 || good) { - good = node.widgets![0]!.value === filename; - if (good) break; - await env.wait(100); - } - if (!good) { - throw new Error("Expected file not loaded."); - } - return node; -} - describe("TestImageOrLatentSize", async () => { await beforeEach(async () => { await env.clear(); @@ -61,21 +30,21 @@ describe("TestImageOrLatentSize", async () => { }); await should("get correct size for a 1x1 image", async () => { - await pastImageToLoadImgeNode(PNG_1x1, imageNode); + await pasteImageToLoadImageNode(env, PNG_1x1, imageNode); await env.queuePrompt(); expect(displayAnyW.widgets![0]!.value).toBe("width", 1); expect(displayAnyH.widgets![0]!.value).toBe("height", 1); }); await should("get correct size for a 1x2 image", async () => { - await pastImageToLoadImgeNode(PNG_1x2, imageNode); + await pasteImageToLoadImageNode(env, PNG_1x2, imageNode); await env.queuePrompt(); expect(displayAnyW.widgets![0]!.value).toBe("width", 1); expect(displayAnyH.widgets![0]!.value).toBe("height", 2); }); await should("get correct size for a 2x1 image", async () => { - await pastImageToLoadImgeNode(PNG_2x1, imageNode); + await pasteImageToLoadImageNode(env, PNG_2x1, imageNode); await env.queuePrompt(); expect(displayAnyW.widgets![0]!.value).toBe("width", 2); expect(displayAnyH.widgets![0]!.value).toBe("height", 1); diff --git a/src_web/comfyui/tests/power_puter.ts b/src_web/comfyui/tests/power_puter.ts index 2edf0ff..165f977 100644 --- a/src_web/comfyui/tests/power_puter.ts +++ b/src_web/comfyui/tests/power_puter.ts @@ -3,16 +3,17 @@ import type {LGraphNode} from "@comfyorg/frontend"; import {NodeTypesString} from "../constants"; import {ComfyUITestEnvironment} from "../testing/comfyui_env"; import {describe, should, beforeEach, expect, describeRun} from "../testing/runner.js"; +import {pasteImageToLoadImageNode, PNG_1x1, PNG_1x2, PNG_2x1} from "../testing/utils_test.js"; const env = new ComfyUITestEnvironment(); function setPowerPuterValue(node: LGraphNode, outputType: string, value: string) { // Strip as much whitespace on first non-empty line from all lines. - if (value.includes('\n')) { - value = value.replace(/^\n/gm, '') + if (value.includes("\n")) { + value = value.replace(/^\n/gm, ""); const strip = value.match(/^(.*?)\S/)?.[1]?.length; if (strip) { - value = value.replace(new RegExp(`^.{${strip}}`, 'mg'), '') + value = value.replace(new RegExp(`^.{${strip}}`, "mg"), ""); } } node.widgets![1]!.value = value; @@ -107,4 +108,23 @@ describe("TestPowerPuter", async () => { await env.queuePrompt(); expect(displayAny.widgets![0]!.value).toBe("(4, 2, 0)"); }); + + await should("disallow calls to some methods", async () => { + const imageNode = await pasteImageToLoadImageNode(env); + imageNode.connect(0, powerPuter, 0); + setPowerPuterValue( + powerPuter, + "STRING", + `a.numpy().tofile('/tmp/test') + `, + ); + await env.queuePrompt(); + + // Check to see if there's an error. + expect(document.querySelector(".p-dialog-mask .p-card-body")!.textContent).toContain( + "error message", + "Disallowed access to \"tofile\" for type ", + ); + (document.querySelector(".p-dialog-mask .p-dialog-close-button")! as HTMLButtonElement).click(); + }); });