WildAi
884c51c962
The pack owned a private stdlib handler, so every line bypassed the coloured console of ComfyUI, and the call sites carried a literal KEEP prefix. The logger now adopts the host console (no handler and propagate=True when ComfyUI is present, one private handler when it is not), the tag lives in a single filter, and the console writes of the pack (the tqdm bars of the processor, the prints of the vendored facelib) reach it through a thread-scoped redirect.
355 lines
12 KiB
Python
355 lines
12 KiB
Python
"""The console routing of the pack - every console write reaches the logger.
|
|
|
|
The vendored facelib writes to the console with ``print()``. Those writes
|
|
bypass the console of ComfyUI: they are untagged and they never reach the
|
|
log file of the host. ``src/keep_console.py`` routes them to the logger of
|
|
the pack:
|
|
|
|
* the writes of the **owning thread** become log records - one record per
|
|
complete line - and the tag filter adds the ``[KEEP]`` tag;
|
|
* the writes of **another thread** are forwarded to the original stream
|
|
untouched;
|
|
* the streams are restored when the block ends - also after an exception;
|
|
* the output of ComfyUI is not swallowed: its root handler holds the stream
|
|
object it was built with and writes straight through the swap.
|
|
|
|
The file also carries the guards of the **tqdm removal**: the processor must
|
|
not import ``tqdm`` any more - the progress of the node is reported through
|
|
the ComfyUI execution API (``src/keep_progress.py``).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import contextlib
|
|
import importlib
|
|
import io
|
|
import logging
|
|
import sys
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parent.parent
|
|
SRC_DIR = REPO_ROOT / "src"
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def console_module():
|
|
return importlib.import_module("ComfyUI_KEEP.src.keep_console")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def logger_module():
|
|
return importlib.import_module("ComfyUI_KEEP.src.keep_logger")
|
|
|
|
|
|
class _Capture(logging.Handler):
|
|
"""A handler that keeps the records it receives."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.records: list[logging.LogRecord] = []
|
|
|
|
def emit(self, record: logging.LogRecord) -> None:
|
|
self.records.append(record)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _records(target: logging.Logger):
|
|
"""The records the logger receives while the block runs."""
|
|
handler = _Capture()
|
|
target.addHandler(handler)
|
|
try:
|
|
yield handler.records
|
|
finally:
|
|
target.removeHandler(handler)
|
|
|
|
|
|
def _pack_logger(logger_module, name: str) -> logging.Logger:
|
|
"""A **dedicated** pack logger - never the global one.
|
|
|
|
The name carries no dot on purpose: a dotted name would make the logger
|
|
a child of ``ComfyUI-KEEP`` and the records would travel through the
|
|
global pack logger.
|
|
"""
|
|
return logger_module.configure_logger(logging.getLogger(name), comfyui_present=True)
|
|
|
|
|
|
def _messages(records) -> list[str]:
|
|
"""The messages of the records."""
|
|
return [record.getMessage() for record in records]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# the capture
|
|
# ---------------------------------------------------------------------------
|
|
@pytest.mark.unit
|
|
def test_a_print_inside_the_block_becomes_a_tagged_record(
|
|
console_module, logger_module
|
|
):
|
|
"""The console write of the owning thread reaches the logger."""
|
|
logger = _pack_logger(logger_module, "KEEP-test-console-print")
|
|
|
|
with _records(logger) as records:
|
|
with console_module.capture_console_to_logger(logger):
|
|
print("hello from the vendored tree")
|
|
|
|
assert _messages(records) == ["[KEEP] hello from the vendored tree"]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_default_logger_is_the_logger_of_the_pack(console_module, logger_module):
|
|
"""Without an argument the pack routes to the logger of the pack."""
|
|
assert console_module.logger is logger_module.logger
|
|
|
|
with _records(logger_module.logger) as records:
|
|
with console_module.capture_console_to_logger():
|
|
print("hello")
|
|
|
|
assert _messages(records) == ["[KEEP] hello"]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_a_partial_write_is_buffered_until_the_line_is_complete(
|
|
console_module, logger_module
|
|
):
|
|
"""One record per **complete** line - a partial line waits for its newline."""
|
|
logger = _pack_logger(logger_module, "KEEP-test-console-partial")
|
|
|
|
with _records(logger) as records:
|
|
with console_module.capture_console_to_logger(logger):
|
|
sys.stdout.write("first ")
|
|
assert _messages(records) == [], "an incomplete line must not be logged"
|
|
sys.stdout.write("line\n")
|
|
sys.stdout.write("second line\n")
|
|
|
|
assert _messages(records) == ["[KEEP] first line", "[KEEP] second line"]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_a_line_without_a_newline_is_flushed_at_the_end_of_the_block(
|
|
console_module, logger_module
|
|
):
|
|
"""The output of a ``print(..., end="")`` is not lost."""
|
|
logger = _pack_logger(logger_module, "KEEP-test-console-flush")
|
|
|
|
with _records(logger) as records:
|
|
with console_module.capture_console_to_logger(logger):
|
|
sys.stdout.write("no newline")
|
|
|
|
assert _messages(records) == ["[KEEP] no newline"]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_stderr_of_the_owning_thread_is_captured_too(
|
|
console_module, logger_module
|
|
):
|
|
"""``sys.stderr`` is captured as well - as a warning."""
|
|
logger = _pack_logger(logger_module, "KEEP-test-console-stderr")
|
|
|
|
with _records(logger) as records:
|
|
with console_module.capture_console_to_logger(logger):
|
|
print("on stderr", file=sys.stderr)
|
|
|
|
assert len(records) == 1
|
|
assert records[0].getMessage() == "[KEEP] on stderr"
|
|
assert records[0].levelno == logging.WARNING
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# the thread scope
|
|
# ---------------------------------------------------------------------------
|
|
@pytest.mark.unit
|
|
def test_the_write_of_another_thread_is_not_captured(
|
|
console_module, logger_module, monkeypatch
|
|
):
|
|
"""Another thread keeps the console - its write is forwarded untouched."""
|
|
logger = _pack_logger(logger_module, "KEEP-test-console-thread")
|
|
original = io.StringIO()
|
|
monkeypatch.setattr(sys, "stdout", original)
|
|
|
|
inside = threading.Event()
|
|
done = threading.Event()
|
|
|
|
def _worker():
|
|
inside.wait(5)
|
|
print("from the worker")
|
|
done.set()
|
|
|
|
with _records(logger) as records:
|
|
with console_module.capture_console_to_logger(logger):
|
|
worker = threading.Thread(target=_worker)
|
|
worker.start()
|
|
inside.set()
|
|
assert done.wait(5), "the worker did not finish"
|
|
worker.join(5)
|
|
|
|
assert _messages(records) == [], "the write of the worker must not be captured"
|
|
assert original.getvalue() == "from the worker\n"
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_two_threads_capture_into_their_own_logger(
|
|
console_module, logger_module
|
|
):
|
|
"""Two capture blocks may run at the same time - each one owns its logger."""
|
|
logger_a = _pack_logger(logger_module, "KEEP-test-console-parallel-a")
|
|
logger_b = _pack_logger(logger_module, "KEEP-test-console-parallel-b")
|
|
|
|
inside_a = threading.Event()
|
|
inside_b = threading.Event()
|
|
errors: list[BaseException] = []
|
|
|
|
def _run(name, target, mine, other):
|
|
try:
|
|
with console_module.capture_console_to_logger(target):
|
|
mine.set()
|
|
other.wait(5)
|
|
print(f"from {name}")
|
|
except BaseException as exc: # noqa: BLE001 - the failure is asserted below
|
|
errors.append(exc)
|
|
|
|
with _records(logger_a) as records_a, _records(logger_b) as records_b:
|
|
thread_a = threading.Thread(target=_run, args=("a", logger_a, inside_a, inside_b))
|
|
thread_a.start()
|
|
assert inside_a.wait(5), "the first thread did not enter the block"
|
|
thread_b = threading.Thread(target=_run, args=("b", logger_b, inside_b, inside_a))
|
|
thread_b.start()
|
|
thread_a.join(5)
|
|
thread_b.join(5)
|
|
|
|
assert errors == []
|
|
assert _messages(records_a) == ["[KEEP] from a"]
|
|
assert _messages(records_b) == ["[KEEP] from b"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# the restoration
|
|
# ---------------------------------------------------------------------------
|
|
@pytest.mark.unit
|
|
def test_the_streams_are_restored_after_the_block(console_module, logger_module):
|
|
"""The block leaves the streams of the process as it found them."""
|
|
logger = _pack_logger(logger_module, "KEEP-test-console-restore")
|
|
before = (sys.stdout, sys.stderr)
|
|
|
|
with console_module.capture_console_to_logger(logger):
|
|
assert (sys.stdout, sys.stderr) != before
|
|
|
|
assert (sys.stdout, sys.stderr) == before
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_streams_are_restored_after_an_exception(
|
|
console_module, logger_module
|
|
):
|
|
"""An exception inside the block does not leave the streams swapped."""
|
|
logger = _pack_logger(logger_module, "KEEP-test-console-restore-error")
|
|
before = (sys.stdout, sys.stderr)
|
|
|
|
with pytest.raises(RuntimeError):
|
|
with console_module.capture_console_to_logger(logger):
|
|
raise RuntimeError("boom")
|
|
|
|
assert (sys.stdout, sys.stderr) == before
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_a_nested_block_restores_the_outer_one(console_module, logger_module):
|
|
"""The blocks may be nested - the inner one does not break the outer."""
|
|
logger = _pack_logger(logger_module, "KEEP-test-console-nested")
|
|
outer = _pack_logger(logger_module, "KEEP-test-console-nested-outer")
|
|
before = (sys.stdout, sys.stderr)
|
|
|
|
with _records(logger) as records_inner, _records(outer) as records_outer:
|
|
with console_module.capture_console_to_logger(outer):
|
|
with console_module.capture_console_to_logger(logger):
|
|
print("inner")
|
|
print("outer")
|
|
|
|
assert (sys.stdout, sys.stderr) == before
|
|
assert _messages(records_outer) == ["[KEEP] outer"]
|
|
assert _messages(records_inner) == ["[KEEP] inner"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# the output of the host
|
|
# ---------------------------------------------------------------------------
|
|
@pytest.mark.unit
|
|
def test_the_output_of_comfyui_is_not_swallowed(
|
|
console_module, logger_module, monkeypatch
|
|
):
|
|
"""The root handler of the host holds its own stream - it is not routed.
|
|
|
|
``app.logger.setup_logger()`` installs the ``LogInterceptor`` before it
|
|
builds the root handler, so the handler writes straight through the swap.
|
|
"""
|
|
logger = _pack_logger(logger_module, "KEEP-test-console-comfy")
|
|
original = io.StringIO()
|
|
monkeypatch.setattr(sys, "stdout", original)
|
|
monkeypatch.setattr(sys, "stderr", original)
|
|
|
|
handler = logging.StreamHandler(original)
|
|
host_logger = logging.getLogger("KEEP-test-comfy-host")
|
|
host_logger.setLevel(logging.WARNING)
|
|
host_logger.addHandler(handler)
|
|
try:
|
|
with console_module.capture_console_to_logger(logger):
|
|
host_logger.warning("comfy says hello")
|
|
finally:
|
|
host_logger.removeHandler(handler)
|
|
|
|
assert "comfy says hello" in original.getvalue()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# the tqdm removal
|
|
# ---------------------------------------------------------------------------
|
|
def _imported_modules(path: Path) -> set[str]:
|
|
"""The top level modules a source file imports."""
|
|
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
|
|
imported: set[str] = set()
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Import):
|
|
imported.update(alias.name.split(".")[0] for alias in node.names)
|
|
elif isinstance(node, ast.ImportFrom) and node.module:
|
|
imported.add(node.module.split(".")[0])
|
|
return imported
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_processor_does_not_import_tqdm():
|
|
"""The progress bars of the processor are gone - the pack logs its steps."""
|
|
source = SRC_DIR / "keep_processor.py"
|
|
assert "tqdm" not in _imported_modules(source)
|
|
assert "tqdm(" not in source.read_text(encoding="utf-8")
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_processor_logs_the_four_steps():
|
|
"""The four steps of the sequence path are logged - one message each."""
|
|
source = (SRC_DIR / "keep_processor.py").read_text(encoding="utf-8")
|
|
tree = ast.parse(source, filename=str(SRC_DIR / "keep_processor.py"))
|
|
|
|
calls = [
|
|
node
|
|
for node in ast.walk(tree)
|
|
if isinstance(node, ast.Call)
|
|
and isinstance(node.func, ast.Attribute)
|
|
and node.func.attr == "info"
|
|
and isinstance(node.func.value, ast.Name)
|
|
and node.func.value.id == "logger"
|
|
and node.args
|
|
and isinstance(node.args[0], ast.Constant)
|
|
and isinstance(node.args[0].value, str)
|
|
]
|
|
# the calls in the order of the source
|
|
messages = [node.args[0].value for node in sorted(calls, key=lambda call: call.lineno)]
|
|
assert messages == [
|
|
"Detecting face landmarks",
|
|
"Cropping and aligning faces",
|
|
"Restoring faces with KEEP",
|
|
"Pasting faces and finalizing frames",
|
|
]
|