222 lines
7.5 KiB
Python
222 lines
7.5 KiB
Python
"""Regression tests for optional v3 input ports that are simply absent.
|
|
|
|
ComfyUI builds node arguments from the keys present in the prompt
|
|
(execution.get_input_data iterates prompt[node]["inputs"]). An unconnected
|
|
optional port (e.g. Save Image's image input) is never serialized, so the
|
|
keyword does not reach execute() at all. Without a Python default the call
|
|
raises "TypeError: ... missing 1 required positional argument" before the
|
|
node's own "if image is None" guard can run.
|
|
"""
|
|
|
|
import ast
|
|
import asyncio
|
|
import inspect
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from _node_loader import load_node
|
|
|
|
ROOT = Path(__file__).parents[1]
|
|
|
|
|
|
def _port(name):
|
|
class _Port:
|
|
def Input(self, *args, **kwargs):
|
|
return ("input", name, args, kwargs)
|
|
|
|
def Output(self, *args, **kwargs):
|
|
return ("output", name, args, kwargs)
|
|
|
|
return _Port()
|
|
|
|
|
|
class _NodeOutput:
|
|
def __init__(self, *values, ui=None, **kwargs):
|
|
self.values = values
|
|
self.ui = ui
|
|
|
|
|
|
def _load_save_image(directory):
|
|
folders = types.ModuleType("folder_paths")
|
|
folders.get_output_directory = lambda: directory
|
|
folders.get_temp_directory = lambda: directory
|
|
|
|
saved_calls = []
|
|
|
|
class _Helper:
|
|
@staticmethod
|
|
def save_images(image, *, filename_prefix, folder_type, cls, compress_level):
|
|
saved_calls.append(
|
|
{
|
|
"shape": tuple(image.shape),
|
|
"filename_prefix": filename_prefix,
|
|
"folder_type": folder_type,
|
|
}
|
|
)
|
|
return [types.SimpleNamespace(subfolder="", filename="ComfyUI_00001_.png")]
|
|
|
|
io = types.SimpleNamespace(
|
|
ComfyNode=object,
|
|
Schema=lambda **kwargs: kwargs,
|
|
NodeOutput=_NodeOutput,
|
|
FolderType=types.SimpleNamespace(output="output", temp="temp"),
|
|
Image=_port("IMAGE"),
|
|
String=_port("STRING"),
|
|
Boolean=_port("BOOLEAN"),
|
|
)
|
|
ui = types.SimpleNamespace(
|
|
ImageSaveHelper=_Helper,
|
|
SavedImages=lambda results: ("saved", results),
|
|
)
|
|
latest = types.ModuleType("comfy_api.latest")
|
|
latest.io = io
|
|
latest.ui = ui
|
|
|
|
previous = {
|
|
name: sys.modules.get(name)
|
|
for name in ("folder_paths", "comfy_api", "comfy_api.latest")
|
|
}
|
|
sys.modules["folder_paths"] = folders
|
|
sys.modules.setdefault("comfy_api", types.ModuleType("comfy_api"))
|
|
sys.modules["comfy_api.latest"] = latest
|
|
try:
|
|
module = load_node("save_image_optional_test", "nodes/io/save_image.py")
|
|
finally:
|
|
for name, value in previous.items():
|
|
if value is None:
|
|
sys.modules.pop(name, None)
|
|
else:
|
|
sys.modules[name] = value
|
|
return module, saved_calls
|
|
|
|
|
|
class SaveImageOptionalInputTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp.cleanup)
|
|
self.module, self.saved_calls = _load_save_image(self.temp.name)
|
|
|
|
def test_execute_reports_default_for_image(self):
|
|
signature = inspect.signature(self.module.SaveImage.execute)
|
|
self.assertIs(signature.parameters["image"].default, None)
|
|
|
|
def test_unconnected_image_is_not_forwarded(self):
|
|
# Mirrors the prompt ComfyUI executes when the port is empty: widget
|
|
# values only, no "image" key at all.
|
|
result = asyncio.run(
|
|
self.module.SaveImage.execute(
|
|
filename="image/ComfyUI",
|
|
auto_increment=True,
|
|
save_output=True,
|
|
save_metadata=True,
|
|
)
|
|
)
|
|
|
|
self.assertEqual(result.values, ())
|
|
self.assertEqual(self.saved_calls, [])
|
|
|
|
def test_bare_execute_call_does_not_raise(self):
|
|
result = asyncio.run(self.module.SaveImage.execute())
|
|
self.assertEqual(result.values, ())
|
|
|
|
def test_connected_image_still_saves(self):
|
|
image = torch.zeros((2, 8, 8, 3), dtype=torch.float32)
|
|
result = asyncio.run(
|
|
self.module.SaveImage.execute(
|
|
image=image,
|
|
filename="image/ComfyUI",
|
|
auto_increment=True,
|
|
save_output=False,
|
|
save_metadata=False,
|
|
)
|
|
)
|
|
|
|
self.assertEqual(result.values, ("",))
|
|
self.assertEqual(
|
|
result.ui,
|
|
("saved", [types.SimpleNamespace(subfolder="", filename="ComfyUI_00001_.png")]),
|
|
)
|
|
self.assertEqual(self.saved_calls[0]["shape"], (2, 8, 8, 3))
|
|
|
|
|
|
_FUNCTION_DEFS = (ast.FunctionDef, ast.AsyncFunctionDef)
|
|
|
|
|
|
def _execute_signature(class_node):
|
|
for statement in class_node.body:
|
|
if isinstance(statement, _FUNCTION_DEFS) and statement.name == "execute":
|
|
return statement
|
|
return None
|
|
|
|
|
|
def _optional_ports_without_schema_default(class_node):
|
|
"""Port-only optional inputs: optional=True and no widget default."""
|
|
ports = set()
|
|
for statement in class_node.body:
|
|
if not isinstance(statement, _FUNCTION_DEFS) or statement.name != "define_schema":
|
|
continue
|
|
for call in ast.walk(statement):
|
|
if not isinstance(call, ast.Call):
|
|
continue
|
|
if not isinstance(call.func, ast.Attribute) or call.func.attr != "Input":
|
|
continue
|
|
keywords = {keyword.arg for keyword in call.keywords}
|
|
if "optional" not in keywords or "default" in keywords:
|
|
continue
|
|
optional = next(
|
|
keyword.value for keyword in call.keywords if keyword.arg == "optional"
|
|
)
|
|
if not (isinstance(optional, ast.Constant) and optional.value is True):
|
|
continue
|
|
if call.args and isinstance(call.args[0], ast.Constant):
|
|
if isinstance(call.args[0].value, str):
|
|
ports.add(call.args[0].value)
|
|
return ports
|
|
|
|
|
|
def _parameters_with_defaults(exec_node):
|
|
arguments = exec_node.args
|
|
positional = arguments.posonlyargs + arguments.args
|
|
defaults = [None] * (len(positional) - len(arguments.defaults)) + list(arguments.defaults)
|
|
values = {
|
|
argument.arg: default is not None
|
|
for argument, default in zip(positional, defaults)
|
|
if argument.arg not in ("cls", "self")
|
|
}
|
|
for argument, default in zip(arguments.kwonlyargs, arguments.kw_defaults):
|
|
values[argument.arg] = default is not None
|
|
return values
|
|
|
|
|
|
class OptionalPortDefaultContractTests(unittest.TestCase):
|
|
"""Any port-only optional input must be defaulted in execute()."""
|
|
|
|
def test_port_only_optional_inputs_have_python_defaults(self):
|
|
offenders = []
|
|
for path in sorted((ROOT / "nodes").rglob("*.py")):
|
|
tree = ast.parse(path.read_text(encoding="utf-8"))
|
|
for class_node in [node for node in ast.walk(tree) if isinstance(node, ast.ClassDef)]:
|
|
ports = _optional_ports_without_schema_default(class_node)
|
|
if not ports:
|
|
continue
|
|
exec_node = _execute_signature(class_node)
|
|
if exec_node is None or exec_node.args.kwarg is not None:
|
|
continue
|
|
defaults = _parameters_with_defaults(exec_node)
|
|
for port in sorted(ports):
|
|
if port not in defaults:
|
|
continue
|
|
if not defaults[port]:
|
|
offenders.append(f"{path.relative_to(ROOT)}::{class_node.name}.{port}")
|
|
|
|
self.assertEqual(offenders, [], "optional ports need a Python default")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|