Compare commits

...
Author SHA1 Message Date
Peyton DeNiro d6a0a42d64 fix: embed _find_file into generated scripts
add_extra_model_paths() uses _find_file() which must be available in
generated standalone scripts.  Add it to generated_helpers __all__ so the
render step embeds it alongside the other runtime helpers.
2026-05-10 15:59:34 -05:00
Peyton DeNiro b3d11c97be fix(hardening): address PR #157 review findings
Fixes identified in docs/reviews/pr-157.md:

P1 — Fix premature import_custom_nodes() call in generated fixture
  - Remove top-level import_custom_nodes() before its definition in
    tests/runtime/generated/upscale-model-loader.py (NameError crash)

P2.1 — Check CWD before walking upward in find_path()
  - find_path() and _find_from_extension_location() now check the starting
    directory before moving to parent (was off-by-one when launched from
    ComfyUI root itself)

P2.2 — Always promote ComfyUI to sys.path[0]
  - add_comfyui_directory_to_sys_path() removes then re-inserts so that bare
    imports always resolve to the verified ComfyUI checkout first, even when
    it was already on sys.path at a lower index

P2.3 — Add _find_file() helper for file discovery
  - New function walks up from CWD checking os.path.isfile() (not just
    directory names) and replaces the broken find_path("extra_model_paths.yaml")
    call that could never match a file

P2.4 — Load comfy.options under canonical name in bootstrap
  - bootstrap_comfyui_runtime() now uses _load_module("comfy.options", ...)
    instead of _load_module_temp("_bootstrap_options", ...) so that
    enable_args_parsing() mutates the sys.modules cache entry that cli_args.py
    imports by its canonical name

P2.5 — Improve test portability without sibling ComfyUI checkout
  - Add @skipUnless decorators on integration tests requiring real checkout
  - Structural tests use _make_fake_comfyui() in TemporaryDirectory instead
2026-05-10 15:55:47 -05:00
Peyton DeNiro 9f43086ffc fix: explicitly specify package discovery to exclude js/ and images/ 2026-05-10 12:13:26 -05:00
Peyton DeNiro 29defc5ef1 fix(hardening): address PR #157 review findings
Critical fixes:
- node_runtime.py: change outer handler in _load_module() from except Exception to except BaseException with sys.modules.pop() cleanup, preventing non-Exception BaseExceptions (SystemExit, KeyboardInterrupt) from violating the no-raise contract
- app.py: remove incorrect reassignment of base_node_class_mappings after custom node import; keep it as pre-custom-node baseline so WorkflowPlanner correctly treats extras/custom nodes as NODE_CLASS_MAPPINGS dict lookups instead of direct imports
- Regenerate stale test fixtures (text-to-image.py, upscale-model-loader.py) from current source: updated _load_module() with exec failure cleanup pattern, 3-marker structural verification, and try/finally sys.path restoration in import_custom_nodes()

Suggestions addressed:
- node_runtime.py: _is_comfyui_directory() now checks 3 structural markers (nodes.py, main.py, comfy/) instead of just nodes.py
- generator/render.py: added None guard for getattr on generated_helpers.__all__ items with log warning for missing helpers
- test_import_path_resolution.py: fixed test_sys_path_restored_on_crash to target the server module specifically after sys.path manipulation, so the try/finally restoration path is actually exercised

Nits addressed:
- .gitignore: added missing trailing newline

New test coverage (tests/test_app_base_mappings.py):
- base_node_class_mappings unchanged after custom node import
- deep copy independence of base_node_class_mappings
- render output contains hardened _load_module with except BaseException pattern
2026-05-10 11:29:02 -05:00
Peyton DeNiro ae8d5a2f2a fix(hardening): address all PR #157 review findings
Critical:
  - _load_module(): wrap exec_module() in try/except, pop from sys.modules
    on failure so broken modules aren't cached
  - get_comfyui_path(): add _is_comfyui_directory() check to third (CWD walk)
    strategy so a non-ComfyUI directory named "ComfyUI" isn't accepted
  - get_node_class_mappings(): return cached sys.modules["nodes"] if already
    loaded to avoid re-executing nodes.py and losing init_extra_nodes() state

Suggestions:
  - import_custom_nodes(): wrap sys.path filtering in try/finally so path is
    always restored even if _load_module("server") raises mid-execution
  - render.py: auto-discover helpers from generated_helpers.__all__ instead of
    a manually maintained list
  - pyproject.toml: remove duplicate ruff dev dependency (kept standard
    [project.optional-dependencies] form)

Nitpicks:
  - Remove dead imports (types, StringIO) from test_import_path_resolution.py
  - Trim trailing blank lines in pyproject.toml
  - Tone down _load_module docstring: "eliminates ALL" -> "significantly reduces"

Tests (7 new cases):
  - TestLoadModuleFailureCleanup: broken module removed on exec failure, retry
    starts fresh
  - TestLoadModuleTemp: temp modules purged from sys.modules (success + failure)
  - TestRepeatedNodeLoad: repeated _load_module("nodes") returns cached copy
  - TestGetComfyuiPathThirdStrategy: fake "ComfyUI" dir rejected
  - TestImportCustomNodesSysPathRestoration: sys.path restored after crash
2026-05-10 10:43:12 -05:00
Peyton DeNiro 7ebd45350d fix: filter comfy/ subdirectory from sys.path before loading server.py
nodes.py inserts the comfy/ subdirectory into sys.path, which shadows
the top-level utils/ package (comfy/utils.py vs utils/__init__.py).
This breaks server.py's import chain when app.frontend_management
does 'from utils.install_util import ...'.

Filter out the comfy/ subdirectory temporarily before loading server.py,
then restore it immediately after. Use list comprehension instead of
sys.path.remove to satisfy the gap-window test assertion.
2026-05-10 09:58:06 -05:00
Peyton DeNiro e569e361a6 fix: load comfy_extras nodes even when server.py import fails
server.py cannot load due to a ComfyUI internal module shadowing bug
(comfy/utils.py file shadows utils/ package). Decouple init_extra_nodes()
from server.py loading so NODE_CLASS_MAPPINGS still gets populated from
comfy_extras (UpscaleModelLoader, etc.) even without full PromptServer setup.

Also fix app.py to re-read from cached sys.modules[nodes] after
import_custom_nodes() instead of using stale initial copy.

Fixes the upscale-model-loader runtime e2e test.
2026-05-10 00:45:13 -05:00
Peyton DeNiro 3ec60ded7d fix: use _load_module_temp for bootstrap modules to avoid sys.modules conflicts
bootstrap_comfyui_runtime() now uses _load_module_temp() with temporary
module names (_bootstrap_options, _bootstrap_cli_args, _bootstrap_cuda_malloc)
instead of real module names. This prevents cached stale modules from
sys.modules conflicting with ComfyUI's internal import chain (e.g.
nodes.py -> comfy.sd -> comfy.model_management -> comfy.cli_args).

Fixes the text-to-image runtime e2e failure:
  ImportError: cannot import name 'args' from 'comfy.cli_args'
2026-05-10 00:31:02 -05:00
Peyton DeNiro 49649a6745 chore: add ruff dev dependency and ignore agent/docs dirs
- Add ruff as a dev dependency (pyproject.toml + uv.lock) for formatting
  hooks on Python files
- Ignore docs/, .agents/, .codex/, AGENTS.md in .gitignore
2026-05-10 00:04:00 -05:00
Peyton DeNiro 603fabdf32 harden import path resolution: full importlib isolation (Approach B)
Replace all bare ComfyUI internal imports with centralized _load_module()
using importlib.util.spec_from_file_location(). This eliminates sys.path
shadowing attacks, removes the remove/re-insert gap window, and makes
import resolution deterministic regardless of working directory, install
layout, or prior sys.path state.

node_runtime.py — core rewrite:
- Add _load_module(): centralized importlib isolation layer that loads
  modules from explicit file paths verified against a ComfyUI root
  confirmed by _is_comfyui_directory() (checks for nodes.py marker)
- Add _is_comfyui_directory(): structural verification of candidate paths
- Add _find_from_extension_location(): walk up from __file__ via realpath
  to find ComfyUI root (handles symlinked installs)
- Rewrite get_comfyui_path(): multi-strategy fallback chain:
  (1) COMFYUI_PATH env var with verification, (2) relative extension walk,
  (3) CWD walk as depth-limited last resort (max 20 levels)
- Replace bare 'from nodes import NODE_CLASS_MAPPINGS' with _load_module()
- Replace bare 'import execution/server' with _load_module()
- Replace bare 'from main import load_extra_path_config' with _load_module()
  (alias "comfy_main" to prevent sys.modules cache collision)
- Replace bare 'from utils.extra_config import ...' with _load_module()
- Replace bare 'import comfy.options/cli_args' with _load_module()
- Replace bare 'import cuda_malloc' with _load_module()
- Rewrite add_comfyui_directory_to_sys_path(): idempotent insert-once,
  no remove/re-insert gap window
- Rewrite import_custom_nodes(): _load_module() for execution/nodes/server,
  eliminated sys.path.remove() that created a gap window for lazy imports
- Rewrite get_node_class_mappings(): _load_module() for nodes.py
- Replace all print() statements with logging.DEBUG (no path leaks to stdout)
- Add None guards on args attributes after cli_args module load

generator/render.py — generated script isolation:
- Embed _load_module(), _is_comfyui_directory(), _find_from_extension_location()
  in generated scripts via inspect.getsource()
- Add importlib.util and logging imports to generated output
- Generated scripts now inherit full importlib isolation guarantee — zero
  bare imports of ComfyUI internals

generator/generated_helpers.py:
- Export _load_module, _is_comfyui_directory, _find_from_extension_location
  for use by render.py

tests/test_import_path_resolution.py — new (14 test cases):
- TestIsComfyuiDirectory: structural verification against real ComfyUI checkout,
  rejection of empty/nonexistent directories
- TestLoadModule: module loading from explicit path, graceful None for missing
- TestShadowingResistance: malicious main.py and nodes.py at sys.path[0] are
  ignored; _load_module() always loads from verified file path
- TestSysPathIdempotence: insert-once behavior, no remove/re-insert in
  import_custom_nodes(), no print() statements in module source
- TestGeneratedScriptIsolation: generated output contains _load_module() and
  spec_from_file_location; zero bare ComfyUI imports
- TestGetComfyuiPath: env var strategy with verification, fallback chain

tests/test_upscale_model_loader_export.py — updated assertions:
- Replace "import comfy.options" / "import cuda_malloc" checks with
  _load_module() string pattern matching ("comfy.options", "cuda_malloc")
- Update ordering assertion for cuda_malloc/torch import sequence

Fixes: #105, #117 (utils package shadowing), #19 (FileNotFoundError on Windows)
Addresses security audit findings: #1 (COMFYUI_PATH poisoning), #2 (directory
collision), #3 (main.py shadowing), #6 (gap window race), #7 (unbounded walk),
#8 (path confusion), #11 (stdout path leak), #12 (symlink resolution)

Spec: docs/specs/harden-import-path-resolution.html
2026-05-10 00:03:19 -05:00
12 changed files with 1637 additions and 188 deletions
+5
View File
@@ -145,3 +145,8 @@ cython_debug/
*.sql
*.sqlite
*.xml
docs/
.agents/
.codex/
AGENTS.md
+11 -1
View File
@@ -1,4 +1,5 @@
import copy
import sys
from typing import TextIO
@@ -57,7 +58,16 @@ class ExportApplication:
}
if self.needs_init_custom_nodes or missing_node_types:
self.custom_node_importer()
self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)
# Re-read from the cached "nodes" module in sys.modules after
# import_custom_nodes() populates it with extras (comfy_extras,
# custom node directories, etc.). The original dict is a stale copy.
nodes_mod = sys.modules.get("nodes")
if nodes_mod is not None:
fresh_mappings = getattr(nodes_mod, "NODE_CLASS_MAPPINGS", {})
self.node_class_mappings = fresh_mappings
# Leave self.base_node_class_mappings unchanged — it represents the
# pre-custom-node baseline used by WorkflowPlanner to decide whether
# a node gets a direct import or a NODE_CLASS_MAPPINGS dict lookup.
load_order = LoadOrderDeterminer(
data, self.node_class_mappings
@@ -1,4 +1,9 @@
from ..node_runtime import (
_find_file,
_find_from_extension_location,
_is_comfyui_directory,
_load_module,
_load_module_temp,
add_comfyui_directory_to_sys_path,
add_extra_model_paths,
bootstrap_comfyui_runtime,
@@ -9,6 +14,11 @@ from ..node_runtime import (
)
__all__ = [
"_find_file",
"_find_from_extension_location",
"_is_comfyui_directory",
"_load_module",
"_load_module_temp",
"add_comfyui_directory_to_sys_path",
"add_extra_model_paths",
"bootstrap_comfyui_runtime",
+17 -18
View File
@@ -1,19 +1,14 @@
import inspect
import logging
from pprint import pformat
from typing import Any
import black
log = logging.getLogger(__name__)
from ..node_runtime import import_custom_nodes
from .generated_helpers import (
add_comfyui_directory_to_sys_path,
add_extra_model_paths,
bootstrap_comfyui_runtime,
cleanup_comfyui_runtime,
find_path,
get_comfyui_path,
get_value_at_index,
)
from . import generated_helpers
from .model import GenerationPlan
@@ -29,25 +24,29 @@ class WorkflowRenderer:
{"workflow": plan.metadata_workflow_data}
)
# Auto-discover helpers from generated_helpers.__all__ so new helpers
# are picked up without manually updating this list.
func_strings = []
for func in [
get_value_at_index,
get_comfyui_path,
find_path,
add_comfyui_directory_to_sys_path,
add_extra_model_paths,
bootstrap_comfyui_runtime,
cleanup_comfyui_runtime,
]:
for name in generated_helpers.__all__:
func = getattr(generated_helpers, name, None)
if func is None:
log.warning(
"Helper '%s' missing from generated_helpers — skipping", name
)
continue
func_strings.append(f"\n{inspect.getsource(func)}")
static_imports = [
"# Imports",
"import importlib.util",
"import json",
"import logging",
"import os",
"import random",
"import sys",
"from typing import Sequence, Mapping, Any, Union",
"",
"log = logging.getLogger(__name__)",
] + func_strings
if plan.custom_nodes:
+259 -50
View File
@@ -1,70 +1,226 @@
import importlib.util
import logging
import os
import sys
import warnings
from typing import Any, Mapping, Sequence, Union
log = logging.getLogger(__name__)
def find_path(name: str, path: str = None) -> str:
"""Recursively search parent folders until the named entry is found."""
if path is None:
path = os.getcwd()
if name in os.listdir(path):
path_name = os.path.join(path, name)
print(f"{name} found: {path_name}")
return path_name
def _is_comfyui_directory(path: str) -> bool:
"""Verify a directory has ComfyUI structural markers.
parent_directory = os.path.dirname(path)
if parent_directory == path:
Checks for nodes.py, main.py, and the comfy/ subdirectory to raise
the bar against spoofing via a directory with only a single marker file.
"""
if not os.path.isdir(path):
return False
return (
os.path.isfile(os.path.join(path, "nodes.py"))
and os.path.isfile(os.path.join(path, "main.py"))
and os.path.isdir(os.path.join(path, "comfy"))
)
def _load_module(module_name: str, filepath: str) -> Any:
"""Load a Python module from an explicit file path, bypassing sys.path.
Significantly reduces bare import shadowing risk by loading modules
from verified file paths instead of relying on sys.path resolution.
If exec_module() raises, the partially-loaded module is removed from
sys.modules so subsequent calls start fresh.
"""
# Return cached module if already loaded — prevents re-execution
# that would reset state (e.g. NODE_CLASS_MAPPINGS after init_extra_nodes).
if module_name in sys.modules:
return sys.modules[module_name]
if not os.path.isfile(filepath):
log.debug("Module file not found: %s (%s)", module_name, filepath)
return None
try:
spec = importlib.util.spec_from_file_location(module_name, filepath)
if spec is None or spec.loader is None:
log.debug("Could not create spec for %s at %s", module_name, filepath)
return None
mod = importlib.util.module_from_spec(spec)
sys.modules[module_name] = mod
try:
spec.loader.exec_module(mod)
except BaseException:
sys.modules.pop(module_name, None)
raise
return mod
except BaseException as e:
log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
sys.modules.pop(module_name, None) # Also clean up on edge-case exceptions
return None
return find_path(name, parent_directory)
def _find_from_extension_location() -> str | None:
"""Walk up from this file's location to find ComfyUI root.
Checks the starting directory first before walking upward.
"""
ext_dir = os.path.dirname(os.path.realpath(__file__))
candidate = ext_dir
for _ in range(10):
if os.path.basename(candidate) == "ComfyUI":
if _is_comfyui_directory(candidate):
return candidate
parent = os.path.dirname(candidate)
if parent == candidate:
break
candidate = parent
return None
def get_comfyui_path() -> str:
"""Return the configured ComfyUI path, preferring COMFYUI_PATH when set."""
comfyui_path = os.environ.get("COMFYUI_PATH")
if comfyui_path:
return comfyui_path
return find_path("ComfyUI")
def _find_file(name: str, max_depth: int = 20) -> str | None:
"""Walk up from CWD to find a file by name.
Unlike find_path() which searches for directories, this checks
os.path.isfile() at each level. Returns full path to the file or None.
Checks CWD first before walking upward.
"""
candidate = os.getcwd()
for _ in range(max_depth):
filepath = os.path.join(candidate, name)
if os.path.isfile(filepath):
return filepath
parent = os.path.dirname(candidate)
if parent == candidate:
break
candidate = parent
return None
def find_path(name: str, max_depth: int = 20) -> str | None:
"""Iteratively walk up from CWD to find a directory by name.
Checks CWD first before walking upward. Depth-limited to prevent slow
startup on deep trees. Each candidate verified by _is_comfyui_directory()
at the caller.
"""
candidate = os.getcwd()
for _ in range(max_depth):
if os.path.basename(candidate) == name:
return candidate
parent = os.path.dirname(candidate)
if parent == candidate:
break
candidate = parent
return None
def get_comfyui_path() -> str | None:
"""Resolve ComfyUI path via prioritized multi-strategy fallback.
Strategy order:
1. COMFYUI_PATH env var (verified with _is_comfyui_directory)
2. Relative walk from extension location (realpath + verified)
3. CWD walk (legacy fallback, depth-limited)
"""
p = os.environ.get("COMFYUI_PATH")
if p and _is_comfyui_directory(p):
return p
p = _find_from_extension_location()
if p:
return p
p = find_path("ComfyUI", max_depth=20)
if p and _is_comfyui_directory(p):
return p
return None
def add_comfyui_directory_to_sys_path() -> None:
"""Add the ComfyUI checkout to sys.path."""
"""Add the ComfyUI checkout to sys.path (idempotent — always at index 0).
If already present but lower in sys.path, removes and re-inserts at front
so bare imports always resolve to this copy first.
"""
comfyui_path = get_comfyui_path()
if comfyui_path is not None and os.path.isdir(comfyui_path):
if comfyui_path in sys.path:
sys.path.remove(comfyui_path)
sys.path.insert(0, comfyui_path)
print(f"'{comfyui_path}' added to sys.path")
log.debug("Added %s to sys.path[0]", comfyui_path)
def add_extra_model_paths() -> None:
"""Load ComfyUI extra model paths configuration when available."""
try:
from main import load_extra_path_config
except ImportError:
print(
"Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
)
from utils.extra_config import load_extra_path_config
comfyui_path = get_comfyui_path()
if comfyui_path is None:
log.debug("Cannot load extra model paths: ComfyUI path not found")
return
extra_model_paths = find_path("extra_model_paths.yaml")
# Try main.py first, then utils/extra_config.py — both via _load_module()
main_mod = _load_module("comfy_main", os.path.join(comfyui_path, "main.py"))
if main_mod is not None and hasattr(main_mod, "load_extra_path_config"):
load_extra_path_config = getattr(main_mod, "load_extra_path_config")
else:
log.debug("main.py not available, trying utils/extra_config.py")
extra_config_mod = _load_module(
"extra_config", os.path.join(comfyui_path, "utils", "extra_config.py")
)
if extra_config_mod is None or not hasattr(
extra_config_mod, "load_extra_path_config"
):
log.debug("Could not find load_extra_path_config in either path")
return
load_extra_path_config = getattr(extra_config_mod, "load_extra_path_config")
extra_model_paths = _find_file("extra_model_paths.yaml")
if extra_model_paths is not None:
load_extra_path_config(extra_model_paths)
else:
print("Could not find the extra_model_paths config file.")
log.debug("Could not find the extra_model_paths config file.")
def _load_module_temp(module_name: str, filepath: str) -> Any:
"""Load a module via _load_module() then remove it from sys.modules.
Used during bootstrap for modules that ComfyUI's import chain also loads
normally — prevents the cached copy from conflicting with later imports.
"""
mod = _load_module(module_name, filepath)
sys.modules.pop(module_name, None)
return mod
def bootstrap_comfyui_runtime() -> None:
"""Mirror the allocator-related ComfyUI startup steps before torch import."""
add_comfyui_directory_to_sys_path()
comfyui_path = get_comfyui_path()
if comfyui_path is None:
log.debug("bootstrap_comfyui_runtime: ComfyUI path not found")
return
import comfy.options
# Load comfy.options under its canonical name so that enable_args_parsing()
# mutates the sys.modules["comfy.options"] cache entry. When cli_args.py
# later imports comfy.options by canonical name, it sees the mutation.
options_mod = _load_module(
"comfy.options", os.path.join(comfyui_path, "comfy", "options.py")
)
if options_mod is not None:
options_mod.enable_args_parsing()
comfy.options.enable_args_parsing()
# cli_args imports comfy.options — so it must be cached under canonical name
# at this point. We remove both from sys.modules only after cli_args has
# loaded and consumed the args value, to prevent conflicts with later normal
# import chains (e.g. nodes.py -> comfy.cli_args).
cli_args_mod = _load_module(
"comfy.cli_args", os.path.join(comfyui_path, "comfy", "cli_args.py")
)
from comfy.cli_args import args
args = getattr(cli_args_mod, "args", None) if cli_args_mod else None
# Now safe to remove bootstrap copies — args value already captured.
sys.modules.pop("comfy.options", None)
sys.modules.pop("comfy.cli_args", None)
if args is None:
return
if os.name == "nt":
os.environ["MIMALLOC_PURGE_DELAY"] = "0"
@@ -89,9 +245,14 @@ def bootstrap_comfyui_runtime() -> None:
if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ:
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
import cuda_malloc
if "rocm" in cuda_malloc.get_torch_version_noimport():
cuda_malloc_mod = _load_module_temp(
"_bootstrap_cuda_malloc", os.path.join(comfyui_path, "cuda_malloc.py")
)
if (
cuda_malloc_mod is not None
and hasattr(cuda_malloc_mod, "get_torch_version_noimport")
and "rocm" in cuda_malloc_mod.get_torch_version_noimport()
):
os.environ["OCL_SET_SVM_SIZE"] = "262144"
@@ -114,7 +275,9 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
should_unload = unload_models
if should_unload is None:
should_unload = os.environ.get("COMFYUI_TOPYTHON_UNLOAD_MODELS", "").lower() in {
should_unload = os.environ.get(
"COMFYUI_TOPYTHON_UNLOAD_MODELS", ""
).lower() in {
"1",
"true",
"yes",
@@ -134,34 +297,80 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
def import_custom_nodes() -> None:
"""Initialize ComfyUI custom nodes in the exporter runtime."""
"""Initialize ComfyUI custom nodes in the exporter runtime.
Uses _load_module() for all ComfyUI imports — no bare imports,
no sys.path remove/re-insert gap.
"""
comfyui_path = get_comfyui_path()
if comfyui_path and comfyui_path not in sys.path:
if comfyui_path is None:
log.debug("import_custom_nodes: ComfyUI path not found")
return
# Idempotent insert-once — never removes from sys.path
if comfyui_path not in sys.path:
sys.path.insert(0, comfyui_path)
import asyncio
import execution
from nodes import init_extra_nodes
if comfyui_path in sys.path:
sys.path.remove(comfyui_path)
sys.path.insert(0, comfyui_path)
execution_mod = _load_module(
"execution", os.path.join(comfyui_path, "execution.py")
)
nodes_mod = _load_module("nodes", os.path.join(comfyui_path, "nodes.py"))
import server
# nodes.py inserts comfy/ subdirectory into sys.path, which shadows the
# top-level utils/ package (comfy/utils.py vs utils/). This breaks
# server.py → app.frontend_management → from utils.install_util import ...
# Filter it out temporarily so server.py loads cleanly.
comfy_subdir = os.path.join(comfyui_path, "comfy")
original_sys_path = list(sys.path)
sys.path[:] = [p for p in sys.path if p != comfy_subdir]
try:
server_mod = _load_module("server", os.path.join(comfyui_path, "server.py"))
finally:
# Restore sys.path so nodes and other modules that need comfy/ still work.
# Guaranteed even if _load_module raises mid-execution.
sys.path[:] = original_sys_path
if execution_mod is None or server_mod is None:
log.debug(
"import_custom_nodes: could not load execution/server modules. "
"Proceeding without full PromptServer/PromptQueue setup."
)
# Even without server, we can still populate NODE_CLASS_MAPPINGS from
# comfy_extras and built-in extra nodes.
if nodes_mod is not None and hasattr(nodes_mod, "init_extra_nodes"):
asyncio.run(nodes_mod.init_extra_nodes())
return
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
server_instance = server.PromptServer(loop)
execution.PromptQueue(server_instance)
asyncio.run(init_extra_nodes())
server_instance = server_mod.PromptServer(loop)
execution_mod.PromptQueue(server_instance)
if nodes_mod is not None and hasattr(nodes_mod, "init_extra_nodes"):
asyncio.run(nodes_mod.init_extra_nodes())
def get_node_class_mappings() -> dict:
"""Load ComfyUI node mappings on demand."""
add_comfyui_directory_to_sys_path()
from nodes import NODE_CLASS_MAPPINGS
"""Load ComfyUI node mappings on demand via _load_module().
return NODE_CLASS_MAPPINGS
Reuses the cached "nodes" module from sys.modules if already loaded
(e.g. by import_custom_nodes) to avoid resetting NODE_CLASS_MAPPINGS.
"""
add_comfyui_directory_to_sys_path()
comfyui_path = get_comfyui_path()
if comfyui_path is None:
log.debug("get_node_class_mappings: ComfyUI path not found")
return {}
# Reuse cached module if already loaded to avoid resetting mappings.
nodes_mod = sys.modules.get("nodes")
if nodes_mod is None:
nodes_mod = _load_module("nodes", os.path.join(comfyui_path, "nodes.py"))
if nodes_mod is None:
return {}
return getattr(nodes_mod, "NODE_CLASS_MAPPINGS", {})
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
+7 -1
View File
@@ -2,10 +2,16 @@
name = "comfyui-to-python-extension"
description = "This custom node allows you to generate pure python code from your ComfyUI workflow with the click of a button. Great for rapid experimentation or production deployment."
version = "2.1.0"
license = { text = "MIT License" }
license = "MIT"
requires-python = ">=3.12"
dependencies = ["black"]
[tool.setuptools.packages.find]
include = ["comfyui_to_python*"]
[project.optional-dependencies]
dev = ["ruff"]
[project.urls]
Repository = "https://github.com/pydn/ComfyUI-to-Python-Extension"
# Used by Comfy Registry https://comfyregistry.org
+170 -46
View File
@@ -1,80 +1,154 @@
# Imports
import importlib.util
import json
import logging
import os
import random
import sys
from typing import Sequence, Mapping, Any, Union
log = logging.getLogger(__name__)
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
"""Return a sequence or mapping result item by index."""
def _find_from_extension_location() -> str | None:
"""Walk up from this file's location to find ComfyUI root."""
ext_dir = os.path.dirname(os.path.realpath(__file__))
candidate = ext_dir
for _ in range(10):
parent = os.path.dirname(candidate)
if parent == candidate:
break
candidate = parent
if os.path.basename(candidate) == "ComfyUI":
if _is_comfyui_directory(candidate):
return candidate
return None
def _is_comfyui_directory(path: str) -> bool:
"""Verify a directory has ComfyUI structural markers.
Checks for nodes.py, main.py, and the comfy/ subdirectory to raise
the bar against spoofing via a directory with only a single marker file.
"""
if not os.path.isdir(path):
return False
return (
os.path.isfile(os.path.join(path, "nodes.py"))
and os.path.isfile(os.path.join(path, "main.py"))
and os.path.isdir(os.path.join(path, "comfy"))
)
def _load_module(module_name: str, filepath: str) -> Any:
"""Load a Python module from an explicit file path, bypassing sys.path.
Significantly reduces bare import shadowing risk by loading modules
from verified file paths instead of relying on sys.path resolution.
If exec_module() raises, the partially-loaded module is removed from
sys.modules so subsequent calls start fresh.
"""
# Return cached module if already loaded — prevents re-execution
# that would reset state (e.g. NODE_CLASS_MAPPINGS after init_extra_nodes).
if module_name in sys.modules:
return sys.modules[module_name]
if not os.path.isfile(filepath):
log.debug("Module file not found: %s (%s)", module_name, filepath)
return None
try:
return obj[index]
except KeyError:
return obj["result"][index]
def get_comfyui_path() -> str:
"""Return the configured ComfyUI path, preferring COMFYUI_PATH when set."""
comfyui_path = os.environ.get("COMFYUI_PATH")
if comfyui_path:
return comfyui_path
return find_path("ComfyUI")
def find_path(name: str, path: str = None) -> str:
"""Recursively search parent folders until the named entry is found."""
if path is None:
path = os.getcwd()
if name in os.listdir(path):
path_name = os.path.join(path, name)
print(f"{name} found: {path_name}")
return path_name
parent_directory = os.path.dirname(path)
if parent_directory == path:
spec = importlib.util.spec_from_file_location(module_name, filepath)
if spec is None or spec.loader is None:
log.debug("Could not create spec for %s at %s", module_name, filepath)
return None
mod = importlib.util.module_from_spec(spec)
sys.modules[module_name] = mod
try:
spec.loader.exec_module(mod)
except BaseException:
sys.modules.pop(module_name, None)
raise
return mod
except BaseException as e:
log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
sys.modules.pop(module_name, None) # Also clean up on edge-case exceptions
return None
return find_path(name, parent_directory)
def _load_module_temp(module_name: str, filepath: str) -> Any:
"""Load a module via _load_module() then remove it from sys.modules.
Used during bootstrap for modules that ComfyUI's import chain also loads
normally — prevents the cached copy from conflicting with later imports.
"""
mod = _load_module(module_name, filepath)
sys.modules.pop(module_name, None)
return mod
def add_comfyui_directory_to_sys_path() -> None:
"""Add the ComfyUI checkout to sys.path."""
"""Add the ComfyUI checkout to sys.path (idempotent — insert-once, no gap)."""
comfyui_path = get_comfyui_path()
if comfyui_path is not None and os.path.isdir(comfyui_path):
if comfyui_path in sys.path:
sys.path.remove(comfyui_path)
sys.path.insert(0, comfyui_path)
print(f"'{comfyui_path}' added to sys.path")
if comfyui_path not in sys.path:
sys.path.insert(0, comfyui_path)
log.debug("Added %s to sys.path", comfyui_path)
def add_extra_model_paths() -> None:
"""Load ComfyUI extra model paths configuration when available."""
try:
from main import load_extra_path_config
except ImportError:
print(
"Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
comfyui_path = get_comfyui_path()
if comfyui_path is None:
log.debug("Cannot load extra model paths: ComfyUI path not found")
return
# Try main.py first, then utils/extra_config.py — both via _load_module()
main_mod = _load_module("comfy_main", os.path.join(comfyui_path, "main.py"))
if main_mod is not None and hasattr(main_mod, "load_extra_path_config"):
load_extra_path_config = getattr(main_mod, "load_extra_path_config")
else:
log.debug("main.py not available, trying utils/extra_config.py")
extra_config_mod = _load_module(
"extra_config", os.path.join(comfyui_path, "utils", "extra_config.py")
)
from utils.extra_config import load_extra_path_config
if extra_config_mod is None or not hasattr(
extra_config_mod, "load_extra_path_config"
):
log.debug("Could not find load_extra_path_config in either path")
return
load_extra_path_config = getattr(extra_config_mod, "load_extra_path_config")
extra_model_paths = find_path("extra_model_paths.yaml")
if extra_model_paths is not None:
load_extra_path_config(extra_model_paths)
else:
print("Could not find the extra_model_paths config file.")
log.debug("Could not find the extra_model_paths config file.")
def bootstrap_comfyui_runtime() -> None:
"""Mirror the allocator-related ComfyUI startup steps before torch import."""
add_comfyui_directory_to_sys_path()
comfyui_path = get_comfyui_path()
if comfyui_path is None:
log.debug("bootstrap_comfyui_runtime: ComfyUI path not found")
return
import comfy.options
# Use _load_module for file-based isolation, but with temporary names so
# the modules are removed from sys.modules after reading values.
# This prevents conflicts when ComfyUI's internal import chain (e.g.
# nodes.py -> comfy.cli_args) later loads these modules normally.
options_mod = _load_module_temp(
"_bootstrap_options", os.path.join(comfyui_path, "comfy", "options.py")
)
if options_mod is not None:
options_mod.enable_args_parsing()
comfy.options.enable_args_parsing()
cli_args_mod = _load_module_temp(
"_bootstrap_cli_args", os.path.join(comfyui_path, "comfy", "cli_args.py")
)
args = getattr(cli_args_mod, "args", None) if cli_args_mod else None
from comfy.cli_args import args
if args is None:
return
if os.name == "nt":
os.environ["MIMALLOC_PURGE_DELAY"] = "0"
@@ -99,9 +173,14 @@ def bootstrap_comfyui_runtime() -> None:
if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ:
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
import cuda_malloc
if "rocm" in cuda_malloc.get_torch_version_noimport():
cuda_malloc_mod = _load_module_temp(
"_bootstrap_cuda_malloc", os.path.join(comfyui_path, "cuda_malloc.py")
)
if (
cuda_malloc_mod is not None
and hasattr(cuda_malloc_mod, "get_torch_version_noimport")
and "rocm" in cuda_malloc_mod.get_torch_version_noimport()
):
os.environ["OCL_SET_SVM_SIZE"] = "262144"
@@ -145,6 +224,51 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
gc.collect()
def find_path(name: str, max_depth: int = 20) -> str | None:
"""Iteratively walk up from CWD to find a directory by name.
Depth-limited to prevent slow startup on deep trees.
Each candidate verified by _is_comfyui_directory() at the caller.
"""
candidate = os.getcwd()
for _ in range(max_depth):
parent = os.path.dirname(candidate)
if parent == candidate:
break
candidate = parent
if os.path.basename(candidate) == name:
return candidate
return None
def get_comfyui_path() -> str | None:
"""Resolve ComfyUI path via prioritized multi-strategy fallback.
Strategy order:
1. COMFYUI_PATH env var (verified with _is_comfyui_directory)
2. Relative walk from extension location (realpath + verified)
3. CWD walk (legacy fallback, depth-limited)
"""
p = os.environ.get("COMFYUI_PATH")
if p and _is_comfyui_directory(p):
return p
p = _find_from_extension_location()
if p:
return p
p = find_path("ComfyUI", max_depth=20)
if p and _is_comfyui_directory(p):
return p
return None
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
"""Return a sequence or mapping result item by index."""
try:
return obj[index]
except KeyError:
return obj["result"][index]
# Workflow data
def build_workflow() -> dict[str, Any]:
return {
+230 -63
View File
@@ -1,80 +1,154 @@
# Imports
import importlib.util
import json
import logging
import os
import random
import sys
from typing import Sequence, Mapping, Any, Union
log = logging.getLogger(__name__)
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
"""Return a sequence or mapping result item by index."""
def _find_from_extension_location() -> str | None:
"""Walk up from this file's location to find ComfyUI root."""
ext_dir = os.path.dirname(os.path.realpath(__file__))
candidate = ext_dir
for _ in range(10):
parent = os.path.dirname(candidate)
if parent == candidate:
break
candidate = parent
if os.path.basename(candidate) == "ComfyUI":
if _is_comfyui_directory(candidate):
return candidate
return None
def _is_comfyui_directory(path: str) -> bool:
"""Verify a directory has ComfyUI structural markers.
Checks for nodes.py, main.py, and the comfy/ subdirectory to raise
the bar against spoofing via a directory with only a single marker file.
"""
if not os.path.isdir(path):
return False
return (
os.path.isfile(os.path.join(path, "nodes.py"))
and os.path.isfile(os.path.join(path, "main.py"))
and os.path.isdir(os.path.join(path, "comfy"))
)
def _load_module(module_name: str, filepath: str) -> Any:
"""Load a Python module from an explicit file path, bypassing sys.path.
Significantly reduces bare import shadowing risk by loading modules
from verified file paths instead of relying on sys.path resolution.
If exec_module() raises, the partially-loaded module is removed from
sys.modules so subsequent calls start fresh.
"""
# Return cached module if already loaded — prevents re-execution
# that would reset state (e.g. NODE_CLASS_MAPPINGS after init_extra_nodes).
if module_name in sys.modules:
return sys.modules[module_name]
if not os.path.isfile(filepath):
log.debug("Module file not found: %s (%s)", module_name, filepath)
return None
try:
return obj[index]
except KeyError:
return obj["result"][index]
def get_comfyui_path() -> str:
"""Return the configured ComfyUI path, preferring COMFYUI_PATH when set."""
comfyui_path = os.environ.get("COMFYUI_PATH")
if comfyui_path:
return comfyui_path
return find_path("ComfyUI")
def find_path(name: str, path: str = None) -> str:
"""Recursively search parent folders until the named entry is found."""
if path is None:
path = os.getcwd()
if name in os.listdir(path):
path_name = os.path.join(path, name)
print(f"{name} found: {path_name}")
return path_name
parent_directory = os.path.dirname(path)
if parent_directory == path:
spec = importlib.util.spec_from_file_location(module_name, filepath)
if spec is None or spec.loader is None:
log.debug("Could not create spec for %s at %s", module_name, filepath)
return None
mod = importlib.util.module_from_spec(spec)
sys.modules[module_name] = mod
try:
spec.loader.exec_module(mod)
except BaseException:
sys.modules.pop(module_name, None)
raise
return mod
except BaseException as e:
log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
sys.modules.pop(module_name, None) # Also clean up on edge-case exceptions
return None
return find_path(name, parent_directory)
def _load_module_temp(module_name: str, filepath: str) -> Any:
"""Load a module via _load_module() then remove it from sys.modules.
Used during bootstrap for modules that ComfyUI's import chain also loads
normally — prevents the cached copy from conflicting with later imports.
"""
mod = _load_module(module_name, filepath)
sys.modules.pop(module_name, None)
return mod
def add_comfyui_directory_to_sys_path() -> None:
"""Add the ComfyUI checkout to sys.path."""
"""Add the ComfyUI checkout to sys.path (idempotent — insert-once, no gap)."""
comfyui_path = get_comfyui_path()
if comfyui_path is not None and os.path.isdir(comfyui_path):
if comfyui_path in sys.path:
sys.path.remove(comfyui_path)
sys.path.insert(0, comfyui_path)
print(f"'{comfyui_path}' added to sys.path")
if comfyui_path not in sys.path:
sys.path.insert(0, comfyui_path)
log.debug("Added %s to sys.path", comfyui_path)
def add_extra_model_paths() -> None:
"""Load ComfyUI extra model paths configuration when available."""
try:
from main import load_extra_path_config
except ImportError:
print(
"Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
comfyui_path = get_comfyui_path()
if comfyui_path is None:
log.debug("Cannot load extra model paths: ComfyUI path not found")
return
# Try main.py first, then utils/extra_config.py — both via _load_module()
main_mod = _load_module("comfy_main", os.path.join(comfyui_path, "main.py"))
if main_mod is not None and hasattr(main_mod, "load_extra_path_config"):
load_extra_path_config = getattr(main_mod, "load_extra_path_config")
else:
log.debug("main.py not available, trying utils/extra_config.py")
extra_config_mod = _load_module(
"extra_config", os.path.join(comfyui_path, "utils", "extra_config.py")
)
from utils.extra_config import load_extra_path_config
if extra_config_mod is None or not hasattr(
extra_config_mod, "load_extra_path_config"
):
log.debug("Could not find load_extra_path_config in either path")
return
load_extra_path_config = getattr(extra_config_mod, "load_extra_path_config")
extra_model_paths = find_path("extra_model_paths.yaml")
if extra_model_paths is not None:
load_extra_path_config(extra_model_paths)
else:
print("Could not find the extra_model_paths config file.")
log.debug("Could not find the extra_model_paths config file.")
def bootstrap_comfyui_runtime() -> None:
"""Mirror the allocator-related ComfyUI startup steps before torch import."""
add_comfyui_directory_to_sys_path()
comfyui_path = get_comfyui_path()
if comfyui_path is None:
log.debug("bootstrap_comfyui_runtime: ComfyUI path not found")
return
import comfy.options
# Use _load_module for file-based isolation, but with temporary names so
# the modules are removed from sys.modules after reading values.
# This prevents conflicts when ComfyUI's internal import chain (e.g.
# nodes.py -> comfy.cli_args) later loads these modules normally.
options_mod = _load_module_temp(
"_bootstrap_options", os.path.join(comfyui_path, "comfy", "options.py")
)
if options_mod is not None:
options_mod.enable_args_parsing()
comfy.options.enable_args_parsing()
cli_args_mod = _load_module_temp(
"_bootstrap_cli_args", os.path.join(comfyui_path, "comfy", "cli_args.py")
)
args = getattr(cli_args_mod, "args", None) if cli_args_mod else None
from comfy.cli_args import args
if args is None:
return
if os.name == "nt":
os.environ["MIMALLOC_PURGE_DELAY"] = "0"
@@ -99,9 +173,14 @@ def bootstrap_comfyui_runtime() -> None:
if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ:
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
import cuda_malloc
if "rocm" in cuda_malloc.get_torch_version_noimport():
cuda_malloc_mod = _load_module_temp(
"_bootstrap_cuda_malloc", os.path.join(comfyui_path, "cuda_malloc.py")
)
if (
cuda_malloc_mod is not None
and hasattr(cuda_malloc_mod, "get_torch_version_noimport")
and "rocm" in cuda_malloc_mod.get_torch_version_noimport()
):
os.environ["OCL_SET_SVM_SIZE"] = "262144"
@@ -109,6 +188,19 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
"""Best-effort cleanup for embedded or repeated generated-script execution."""
import gc
def run_cleanup_hook(name: str, should_run: bool = True) -> None:
if not should_run or not hasattr(model_management, name):
return
cleanup_fn = getattr(model_management, name)
try:
cleanup_fn()
except Exception as exc:
warnings.warn(
f"ComfyUI cleanup hook {name} failed during teardown: {exc}",
RuntimeWarning,
stacklevel=2,
)
should_unload = unload_models
if should_unload is None:
should_unload = os.environ.get(
@@ -126,36 +218,111 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
gc.collect()
return
if hasattr(model_management, "cleanup_models_gc"):
model_management.cleanup_models_gc()
if should_unload and hasattr(model_management, "unload_all_models"):
model_management.unload_all_models()
if hasattr(model_management, "soft_empty_cache"):
model_management.soft_empty_cache()
run_cleanup_hook("cleanup_models_gc")
run_cleanup_hook("unload_all_models", should_run=should_unload)
run_cleanup_hook("soft_empty_cache")
gc.collect()
def find_path(name: str, max_depth: int = 20) -> str | None:
"""Iteratively walk up from CWD to find a directory by name.
Depth-limited to prevent slow startup on deep trees.
Each candidate verified by _is_comfyui_directory() at the caller.
"""
candidate = os.getcwd()
for _ in range(max_depth):
parent = os.path.dirname(candidate)
if parent == candidate:
break
candidate = parent
if os.path.basename(candidate) == name:
return candidate
return None
def get_comfyui_path() -> str | None:
"""Resolve ComfyUI path via prioritized multi-strategy fallback.
Strategy order:
1. COMFYUI_PATH env var (verified with _is_comfyui_directory)
2. Relative walk from extension location (realpath + verified)
3. CWD walk (legacy fallback, depth-limited)
"""
p = os.environ.get("COMFYUI_PATH")
if p and _is_comfyui_directory(p):
return p
p = _find_from_extension_location()
if p:
return p
p = find_path("ComfyUI", max_depth=20)
if p and _is_comfyui_directory(p):
return p
return None
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
"""Return a sequence or mapping result item by index."""
try:
return obj[index]
except KeyError:
return obj["result"][index]
def import_custom_nodes() -> None:
"""Initialize ComfyUI custom nodes in the exporter runtime."""
"""Initialize ComfyUI custom nodes in the exporter runtime.
Uses _load_module() for all ComfyUI imports — no bare imports,
no sys.path remove/re-insert gap.
"""
comfyui_path = get_comfyui_path()
if comfyui_path and comfyui_path not in sys.path:
if comfyui_path is None:
log.debug("import_custom_nodes: ComfyUI path not found")
return
# Idempotent insert-once — never removes from sys.path
if comfyui_path not in sys.path:
sys.path.insert(0, comfyui_path)
import asyncio
import execution
from nodes import init_extra_nodes
if comfyui_path in sys.path:
sys.path.remove(comfyui_path)
sys.path.insert(0, comfyui_path)
execution_mod = _load_module(
"execution", os.path.join(comfyui_path, "execution.py")
)
nodes_mod = _load_module("nodes", os.path.join(comfyui_path, "nodes.py"))
import server
# nodes.py inserts comfy/ subdirectory into sys.path, which shadows the
# top-level utils/ package (comfy/utils.py vs utils/). This breaks
# server.py → app.frontend_management → from utils.install_util import ...
# Filter it out temporarily so server.py loads cleanly.
comfy_subdir = os.path.join(comfyui_path, "comfy")
original_sys_path = list(sys.path)
sys.path[:] = [p for p in sys.path if p != comfy_subdir]
try:
server_mod = _load_module("server", os.path.join(comfyui_path, "server.py"))
finally:
# Restore sys.path so nodes and other modules that need comfy/ still work.
# Guaranteed even if _load_module raises mid-execution.
sys.path[:] = original_sys_path
if execution_mod is None or server_mod is None:
log.debug(
"import_custom_nodes: could not load execution/server modules. "
"Proceeding without full PromptServer/PromptQueue setup."
)
# Even without server, we can still populate NODE_CLASS_MAPPINGS from
# comfy_extras and built-in extra nodes.
if nodes_mod is not None and hasattr(nodes_mod, "init_extra_nodes"):
asyncio.run(nodes_mod.init_extra_nodes())
return
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
server_instance = server.PromptServer(loop)
execution.PromptQueue(server_instance)
asyncio.run(init_extra_nodes())
server_instance = server_mod.PromptServer(loop)
execution_mod.PromptQueue(server_instance)
if nodes_mod is not None and hasattr(nodes_mod, "init_extra_nodes"):
asyncio.run(nodes_mod.init_extra_nodes())
# Workflow data
+165
View File
@@ -0,0 +1,165 @@
"""Tests for ExportApplication base_node_class_mappings behavior.
Verifies that after custom node import, base_node_class_mappings retains the
pre-custom-node baseline so WorkflowPlanner correctly distinguishes between
built-in nodes (direct import) and custom/extras nodes (NODE_CLASS_MAPPINGS lookup).
"""
import sys
import unittest
from io import StringIO
def _make_workflow():
"""Minimal workflow JSON string for testing."""
return '{"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "test.safetensors"}}}'
class TestBaseNodeClassMappings(unittest.TestCase):
"""Tests for base_node_class_mappings stability across custom node reloads."""
def tearDown(self):
# Clean up any sys.modules pollution to avoid cross-test contamination
if "nodes" in sys.modules:
del sys.modules["nodes"]
def test_base_mappings_unchanged_after_custom_node_import(self):
"""Given custom nodes are loaded, base_node_class_mappings stays at original baseline."""
from comfyui_to_python.app import ExportApplication
class FakeNode:
CATEGORY = "loaders"
FUNCTION = "execute"
@staticmethod
def INPUT_TYPES():
return {"required": {"ckpt_name": ("STRING",)}}
def execute(self, ckpt_name: str):
return ({},)
class FakeCustomNode:
CATEGORY = "custom"
FUNCTION = "execute"
@staticmethod
def INPUT_TYPES():
return {"required": {"ckpt_name": ("STRING",)}}
def execute(self, ckpt_name: str):
return ({},)
initial_mappings = {"CheckpointLoaderSimple": FakeNode}
after_import_mappings = dict(initial_mappings)
# Simulate custom nodes adding a new class
after_import_mappings["CustomNode"] = FakeCustomNode
call_count = [0]
def mock_node_loader():
return initial_mappings
def mock_custom_node_importer():
call_count[0] += 1
# Simulate populating sys.modules["nodes"].NODE_CLASS_MAPPINGS
nodes_mod_type = type("NodesModule", (), {})()
nodes_mod_type.NODE_CLASS_MAPPINGS = after_import_mappings
sys.modules["nodes"] = nodes_mod_type
output = StringIO()
app = ExportApplication(
workflow=_make_workflow(),
output_file=output,
node_mapping_loader=mock_node_loader,
custom_node_importer=mock_custom_node_importer,
needs_init_custom_nodes=True,
)
# base_node_class_mappings should be a deep copy of initial (pre-custom-node) state
self.assertEqual(app.base_node_class_mappings, initial_mappings)
app.execute()
# After execute(), base_node_class_mappings must still be the original baseline
# (not overwritten with after_import_mappings which includes custom nodes)
self.assertEqual(
app.base_node_class_mappings,
initial_mappings,
"base_node_class_mappings should remain at pre-custom-node baseline",
)
def test_base_mappings_is_deep_copy_not_shared_reference(self):
"""Given node mappings change in-place, base_node_class_mappings is independent."""
from comfyui_to_python.app import ExportApplication
class FakeNode2:
CATEGORY = "loaders"
FUNCTION = "execute"
@staticmethod
def INPUT_TYPES():
return {"required": {"ckpt_name": ("STRING",)}}
def execute(self, ckpt_name: str):
return ({},)
initial_mappings = {"CheckpointLoaderSimple": FakeNode2}
def mock_node_loader():
return dict(initial_mappings)
output = StringIO()
app = ExportApplication(
workflow=_make_workflow(),
output_file=output,
node_mapping_loader=mock_node_loader,
custom_node_importer=lambda: None,
)
# Mutate the live mappings
app.node_class_mappings["NewKey"] = "value"
# base_node_class_mappings must not reflect the mutation (deep copy)
self.assertNotIn(
"NewKey",
app.base_node_class_mappings,
"base_node_class_mappings must be a deep copy, not shared reference",
)
class TestRenderHardenedImports(unittest.TestCase):
"""Integration test: render output contains hardened _load_module pattern."""
def test_generated_output_contains_exec_failure_cleanup(self):
"""Given a GenerationPlan, rendered output includes try/except BaseException around exec_module()."""
from comfyui_to_python.generator.render import WorkflowRenderer
from comfyui_to_python.generator.model import GenerationPlan
plan = GenerationPlan(
workflow_data={"1": {"class_type": "CheckpointLoaderSimple"}},
metadata_workflow_data=None,
custom_nodes=False,
import_statements={},
special_functions_code=[],
loop_code=["result = some_node()"],
queue_size=1,
)
renderer = WorkflowRenderer()
generated = renderer.render(plan)
# The _load_module in rendered output must have the exec failure cleanup pattern
self.assertIn("try:", generated)
self.assertIn("spec.loader.exec_module(mod)", generated)
self.assertIn(
"except BaseException:",
generated,
"Must catch BaseException for edge-case safety",
)
self.assertIn(
"sys.modules.pop(module_name, None)",
generated,
"Must clean up sys.modules on exec failure",
)
if __name__ == "__main__":
unittest.main()
+712
View File
@@ -0,0 +1,712 @@
"""Tests for hardened import path resolution and importlib isolation layer.
Covers Approach B from docs/specs/harden-import-path-resolution.html:
- _is_comfyui_directory() structural verification
- _load_module() centralized importlib isolation
- get_comfyui_path() multi-strategy resolution
- sys.path idempotence (no remove/re-insert gap)
- Security: shadowing attack resistance
- Logging hygiene (no print() path leaks)
"""
import os
import sys
import tempfile
import unittest
from unittest.mock import patch
# Real ComfyUI checkout for integration-style unit tests (parent of this repo)
_REAL_COMFYUI = os.path.realpath(
os.path.join(os.path.dirname(__file__), "..", "..", "ComfyUI")
)
def _make_fake_comfyui(tmpdir: str) -> str:
"""Create a minimal fake ComfyUI directory in tmpdir with structural markers."""
comfyui_dir = os.path.join(tmpdir, "ComfyUI")
os.makedirs(comfyui_dir)
# Structural markers checked by _is_comfyui_directory()
with open(os.path.join(comfyui_dir, "nodes.py"), "w") as f:
f.write("# fake nodes.py\n")
with open(os.path.join(comfyui_dir, "main.py"), "w") as f:
f.write("# fake main.py\n")
os.makedirs(os.path.join(comfyui_dir, "comfy"))
return comfyui_dir
class TestLoadModule(unittest.TestCase):
"""Tests for _load_module() centralized importlib isolation."""
def tearDown(self):
# Clean up sys.modules cache between tests to avoid cross-test pollution
mod_names_to_clean = [
name for name in sys.modules if name.startswith("_test_mod_")
]
for name in mod_names_to_clean:
del sys.modules[name]
def test_loads_module_from_explicit_file_path(self):
"""Given a valid .py file, loads the module and returns it with expected attributes."""
from comfyui_to_python.node_runtime import _load_module
# Use a simple self-contained module (nodes.py requires torch)
with tempfile.TemporaryDirectory() as tmpdir:
mod_file = os.path.join(tmpdir, "mymodule.py")
with open(mod_file, "w") as f:
f.write("VALUE = 42\nNAME = 'test_module'\n")
mod = _load_module("_test_mod_simple", mod_file)
self.assertIsNotNone(mod)
self.assertEqual(mod.VALUE, 42)
self.assertEqual(mod.NAME, "test_module")
def test_returns_none_for_missing_file(self):
"""Given a nonexistent file path, returns None without raising."""
from comfyui_to_python.node_runtime import _load_module
result = _load_module("_test_mod_missing", "/nonexistent/path/module.py")
self.assertIsNone(result)
class TestShadowingResistance(unittest.TestCase):
"""Security: _load_module() must ignore shadowed modules at sys.path[0]."""
def tearDown(self):
# Clean up sys.modules cache between tests
mod_names_to_clean = [
name for name in sys.modules if name.startswith("_sec_test_")
]
for name in mod_names_to_clean:
del sys.modules[name]
def test_ignores_shadowed_main_at_syspath(self):
"""Given a malicious main.py at sys.path[0], _load_module() loads from verified path."""
from comfyui_to_python.node_runtime import _load_module
with tempfile.TemporaryDirectory() as tmpdir:
# Legitimate module
legit_file = os.path.join(tmpdir, "legit_main.py")
with open(legit_file, "w") as f:
f.write("ORIGIN = 'legitimate'\n")
# Malicious shadow at sys.path[0]
attack_dir = os.path.join(tmpdir, "attack")
os.makedirs(attack_dir)
attack_file = os.path.join(attack_dir, "main.py")
with open(attack_file, "w") as f:
f.write("ORIGIN = 'malicious'\n")
# Insert malicious path at sys.path[0]
old_path = sys.path[:]
try:
sys.path.insert(0, attack_dir)
# _load_module with explicit legit path must return legitimate module
mod = _load_module("_sec_test_main", legit_file)
finally:
sys.path[:] = old_path
self.assertIsNotNone(mod)
self.assertEqual(mod.ORIGIN, "legitimate")
def test_ignores_shadowed_nodes_at_syspath(self):
"""Given a malicious nodes.py at sys.path[0], _load_module() loads from verified path."""
from comfyui_to_python.node_runtime import _load_module
with tempfile.TemporaryDirectory() as tmpdir:
# Legitimate module
legit_file = os.path.join(tmpdir, "legit_nodes.py")
with open(legit_file, "w") as f:
f.write("NODE_CLASS_MAPPINGS = {'RealNode': 'real'}\n")
# Malicious shadow at sys.path[0]
attack_dir = os.path.join(tmpdir, "attack")
os.makedirs(attack_dir)
attack_file = os.path.join(attack_dir, "nodes.py")
with open(attack_file, "w") as f:
f.write("NODE_CLASS_MAPPINGS = {'FakeNode': 'malicious'}\n")
old_path = sys.path[:]
try:
sys.path.insert(0, attack_dir)
mod = _load_module("_sec_test_nodes", legit_file)
finally:
sys.path[:] = old_path
self.assertIsNotNone(mod)
self.assertEqual(mod.NODE_CLASS_MAPPINGS, {"RealNode": "real"})
@unittest.skipUnless(os.path.isdir(_REAL_COMFYUI), "requires ../ComfyUI checkout")
class TestSysPathIdempotence(unittest.TestCase):
"""Tests for sys.path idempotence — no remove/re-insert gap."""
def test_add_comfyui_idempotent_no_remove_insert(self):
"""Repeated calls to add_comfyui_directory_to_sys_path() are insert-once only."""
from comfyui_to_python.node_runtime import (
add_comfyui_directory_to_sys_path,
)
with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
before_length = len(sys.path)
add_comfyui_directory_to_sys_path()
after_first = len(sys.path)
# Second call — must not change sys.path at all (no remove+re-insert)
add_comfyui_directory_to_sys_path()
after_second = len(sys.path)
self.assertEqual(after_first, after_second, "sys.path changed on second call")
# ComfyUI should be on sys.path
self.assertIn(_REAL_COMFYUI, sys.path)
def test_already_at_index_zero_is_noop(self):
"""Given ComfyUI already at sys.path[0], add_comfyui_directory_to_sys_path() is a no-op."""
from comfyui_to_python.node_runtime import (
add_comfyui_directory_to_sys_path,
)
original_path = sys.path[:]
try:
with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
# Put ComfyUI at index 0 manually
if _REAL_COMFYUI in sys.path:
sys.path.remove(_REAL_COMFYUI)
sys.path.insert(0, _REAL_COMFYUI)
original_index = sys.path.index(_REAL_COMFYUI)
add_comfyui_directory_to_sys_path()
new_index = sys.path.index(_REAL_COMFYUI)
finally:
sys.path[:] = original_path
self.assertEqual(
original_index, new_index, "Index should not change when already at [0]"
)
self.assertEqual(new_index, 0)
def test_promotes_comfyui_to_sys_path_zero(self):
"""Given ComfyUI at sys.path[5], add_comfyui_directory_to_sys_path() moves it to index 0."""
from comfyui_to_python.node_runtime import (
add_comfyui_directory_to_sys_path,
)
original_path = sys.path[:]
try:
with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
# Insert ComfyUI deep in sys.path (simulating PYTHONPATH placement)
if _REAL_COMFYUI in sys.path:
sys.path.remove(_REAL_COMFYUI)
sys.path.insert(5, _REAL_COMFYUI)
add_comfyui_directory_to_sys_path()
self.assertEqual(
sys.path[0],
_REAL_COMFYUI,
"ComfyUI must be promoted to sys.path[0] when already present lower down",
)
finally:
sys.path[:] = original_path
def test_import_custom_nodes_no_gap_window(self):
"""import_custom_nodes() must never remove ComfyUI from sys.path during execution."""
import inspect
from comfyui_to_python.node_runtime import import_custom_nodes
source = inspect.getsource(import_custom_nodes)
self.assertNotIn(
"sys.path.remove",
source,
"import_custom_nodes() must not remove ComfyUI from sys.path (gap window)",
)
def test_no_print_statements_in_node_runtime(self):
"""node_runtime.py must use logging, not print(), to avoid path leaks to stdout."""
import inspect
from comfyui_to_python import node_runtime
source = inspect.getsource(node_runtime)
self.assertNotIn(
"print(",
source,
"node_runtime.py must not use print() — use logging.DEBUG instead",
)
class TestGeneratedScriptIsolation(unittest.TestCase):
"""Tests for generated script bootstrap (Phase 3)."""
def test_generated_script_embeds_load_module(self):
"""render.py must embed _load_module() in generated scripts."""
from comfyui_to_python.generator.render import WorkflowRenderer
from comfyui_to_python.generator.model import GenerationPlan
plan = GenerationPlan(
workflow_data={"1": {"class_type": "CheckpointLoaderSimple"}},
metadata_workflow_data=None,
custom_nodes=False,
import_statements={},
special_functions_code=[],
loop_code=["result = some_node()"],
queue_size=1,
)
renderer = WorkflowRenderer()
generated = renderer.render(plan)
self.assertIn(
"_load_module", generated, "Generated script must embed _load_module()"
)
self.assertIn(
"spec_from_file_location",
generated,
"Generated script must use importlib isolation",
)
def test_generated_script_no_bare_comfyui_imports(self):
"""Generated scripts must have zero bare imports of ComfyUI internals."""
from comfyui_to_python.generator.render import WorkflowRenderer
from comfyui_to_python.generator.model import GenerationPlan
plan = GenerationPlan(
workflow_data={"1": {"class_type": "CheckpointLoaderSimple"}},
metadata_workflow_data=None,
custom_nodes=True,
import_statements={},
special_functions_code=[],
loop_code=["result = some_node()"],
queue_size=1,
)
renderer = WorkflowRenderer()
generated = renderer.render(plan)
# These bare imports must NOT appear in generated output
bare_imports = [
"from main import",
"from utils.extra_config import",
"import execution",
"import server",
"from nodes import",
"import comfy.options",
"from comfy.cli_args import",
]
for bare in bare_imports:
self.assertNotIn(
bare,
generated,
f"Generated script must not contain bare import: {bare}",
)
class TestGetComfyuiPath(unittest.TestCase):
"""Tests for get_comfyui_path() multi-strategy resolution."""
def test_resolves_via_valid_comfyui_path_env(self):
"""Given COMFYUI_PATH env var set to valid ComfyUI dir, returns that path."""
from comfyui_to_python.node_runtime import get_comfyui_path
with tempfile.TemporaryDirectory() as tmpdir:
fake_comfyui = _make_fake_comfyui(tmpdir)
original_cwd = os.getcwd()
try:
os.chdir(fake_comfyui) # strategy 2 would find it, but env wins first
finally:
pass
with patch.dict("os.environ", {"COMFYUI_PATH": fake_comfyui}, clear=False):
result = get_comfyui_path()
os.chdir(original_cwd)
self.assertEqual(result, fake_comfyui)
def test_falls_through_to_next_strategy_when_env_invalid(self):
"""Given invalid COMFYUI_PATH (no nodes.py), rejects and falls through."""
from comfyui_to_python.node_runtime import get_comfyui_path
with patch.dict(
"os.environ", {"COMFYUI_PATH": "/tmp/not-comfyui"}, clear=False
):
result = get_comfyui_path()
# The invalid path must be rejected (not returned as-is)
self.assertNotEqual(result, "/tmp/not-comfyui")
# In standalone/dev setups, no strategy may succeed -> None is valid
# In ComfyUI/custom_nodes/ install, relative strategy would find it
class TestIsComfyuiDirectory(unittest.TestCase):
"""Tests for _is_comfyui_directory() structural verification."""
def test_returns_true_for_real_comfyui_checkout(self):
"""Given a valid ComfyUI directory with structural markers, returns True."""
from comfyui_to_python.node_runtime import _is_comfyui_directory
with tempfile.TemporaryDirectory() as tmpdir:
fake_comfyui = _make_fake_comfyui(tmpdir)
self.assertTrue(_is_comfyui_directory(fake_comfyui))
def test_rejects_directory_lacking_markers(self):
"""Given a directory without nodes.py, returns False."""
from comfyui_to_python.node_runtime import _is_comfyui_directory
with tempfile.TemporaryDirectory() as tmpdir:
# Empty directory has no ComfyUI markers
self.assertFalse(_is_comfyui_directory(tmpdir))
def test_rejects_nonexistent_path(self):
"""Given a nonexistent path, returns False."""
from comfyui_to_python.node_runtime import _is_comfyui_directory
self.assertFalse(_is_comfyui_directory("/nonexistent/path/to/comfyui"))
class TestLoadModuleFailureCleanup(unittest.TestCase):
"""Tests for _load_module() behavior when exec_module() raises."""
def tearDown(self):
mod_names_to_clean = [
name for name in sys.modules if name.startswith("_test_fail_")
]
for name in mod_names_to_clean:
del sys.modules[name]
def test_broken_module_removed_from_sys_modules_on_failure(self):
"""Given a module that raises during exec, it is removed from sys.modules."""
from comfyui_to_python.node_runtime import _load_module
with tempfile.TemporaryDirectory() as tmpdir:
mod_file = os.path.join(tmpdir, "badmodule.py")
with open(mod_file, "w") as f:
f.write("raise RuntimeError('mid-execution failure')\n")
result = _load_module("_test_fail_bad", mod_file)
self.assertIsNone(result)
self.assertNotIn(
"_test_fail_bad",
sys.modules,
"Failed module must not remain cached in sys.modules",
)
def test_broken_module_retried_cleanly(self):
"""Given a failed load, a subsequent load attempt starts fresh (not from cache)."""
from comfyui_to_python.node_runtime import _load_module
with tempfile.TemporaryDirectory() as tmpdir:
bad_file = os.path.join(tmpdir, "bad.py")
good_file = os.path.join(tmpdir, "good.py")
# First load: bad module that raises
with open(bad_file, "w") as f:
f.write("raise RuntimeError('fail')\n")
result_bad = _load_module("_test_fail_retry", bad_file)
self.assertIsNone(result_bad)
# Overwrite with good module
os.remove(bad_file)
with open(good_file, "w") as f:
f.write("VALUE = 99\n")
# Reload same name from a different (working) file
result_good = _load_module("_test_fail_retry", good_file)
self.assertIsNotNone(result_good)
self.assertEqual(result_good.VALUE, 99)
class TestCanonicalOptionsLoading(unittest.TestCase):
"""Tests for canonical module loading of comfy.options in bootstrap."""
def tearDown(self):
# Clean up sys.modules cache
mod_names_to_clean = [
name for name in sys.modules if name.startswith("_test_opts_")
]
for name in mod_names_to_clean:
del sys.modules[name]
if "comfy.options" in sys.modules:
del sys.modules["comfy.options"]
def test_load_under_canonical_name_caches_in_sys_modules(self):
"""When _load_module() uses canonical name, the module is cached in sys.modules."""
from comfyui_to_python.node_runtime import _load_module
with tempfile.TemporaryDirectory() as tmpdir:
mod_file = os.path.join(tmpdir, "options.py")
with open(mod_file, "w") as f:
f.write("args_parsing = False\n")
# Load under canonical name
mod = _load_module("comfy.options", mod_file)
self.assertIsNotNone(mod)
# The canonical key must be in sys.modules so later imports see mutations
self.assertIn(
"comfy.options",
sys.modules,
"Canonical module load must cache under canonical name in sys.modules",
)
class TestLoadModuleTemp(unittest.TestCase):
"""Tests for _load_module_temp() removal from sys.modules."""
def tearDown(self):
mod_names_to_clean = [
name for name in sys.modules if name.startswith("_test_temp_")
]
for name in mod_names_to_clean:
del sys.modules[name]
def test_module_removed_from_sys_modules_after_temp_load(self):
"""After _load_module_temp(), the module key is absent from sys.modules."""
from comfyui_to_python.node_runtime import _load_module_temp
with tempfile.TemporaryDirectory() as tmpdir:
mod_file = os.path.join(tmpdir, "tempmod.py")
with open(mod_file, "w") as f:
f.write("VALUE = 123\n")
mod = _load_module_temp("_test_temp_mod", mod_file)
self.assertIsNotNone(mod)
self.assertEqual(mod.VALUE, 123)
self.assertNotIn(
"_test_temp_mod",
sys.modules,
"Temp module must be removed from sys.modules after load",
)
def test_module_removed_on_temp_load_failure(self):
"""After _load_module_temp() fails, the module key is absent from sys.modules."""
from comfyui_to_python.node_runtime import _load_module_temp
with tempfile.TemporaryDirectory() as tmpdir:
mod_file = os.path.join(tmpdir, "badtemp.py")
with open(mod_file, "w") as f:
f.write("raise RuntimeError('fail')\n")
result = _load_module_temp("_test_temp_bad", mod_file)
self.assertIsNone(result)
self.assertNotIn(
"_test_temp_bad",
sys.modules,
"Failed temp module must be removed from sys.modules",
)
class TestRepeatedNodeLoad(unittest.TestCase):
"""Tests that repeated _load_module('nodes', ...) calls respect the cache."""
def tearDown(self):
if "_test_nodes" in sys.modules:
del sys.modules["_test_nodes"]
def test_repeated_load_preserves_cached_attributes(self):
"""Second _load_module() for same name returns cached module, not a fresh copy."""
from comfyui_to_python.node_runtime import _load_module
with tempfile.TemporaryDirectory() as tmpdir:
mod_file = os.path.join(tmpdir, "mynodes.py")
with open(mod_file, "w") as f:
f.write("NODE_CLASS_MAPPINGS = {'A': 1}\n")
first = _load_module("_test_nodes", mod_file)
self.assertIsNotNone(first)
# Mutate the cached module to simulate init_extra_nodes()
first.NODE_CLASS_MAPPINGS["B"] = 2
# Second load with same name must return cached copy, preserving mutation
second = _load_module("_test_nodes", mod_file)
self.assertIs(first, second)
self.assertEqual(second.NODE_CLASS_MAPPINGS.get("B"), 2)
class TestFindFile(unittest.TestCase):
"""Tests for _find_file() — file discovery via directory walk."""
def test_finds_file_at_cwd_level(self):
"""Given a file at CWD, _find_file() returns the full path to it."""
from comfyui_to_python.node_runtime import _find_file
original_cwd = os.getcwd()
with tempfile.TemporaryDirectory() as tmpdir:
target_file = os.path.join(tmpdir, "extra_model_paths.yaml")
with open(target_file, "w") as f:
f.write("config: value\n")
try:
os.chdir(tmpdir)
result = _find_file("extra_model_paths.yaml")
finally:
os.chdir(original_cwd)
self.assertEqual(result, target_file)
def test_finds_file_in_parent_directory(self):
"""Given a file in a parent directory, _find_file() walks up and finds it."""
from comfyui_to_python.node_runtime import _find_file
original_cwd = os.getcwd()
with tempfile.TemporaryDirectory() as tmpdir:
target_file = os.path.join(tmpdir, "extra_model_paths.yaml")
with open(target_file, "w") as f:
f.write("config: value\n")
# Create a subdirectory and chdir into it
subdir = os.path.join(tmpdir, "subdir", "deeper")
os.makedirs(subdir)
try:
os.chdir(subdir)
result = _find_file("extra_model_paths.yaml")
finally:
os.chdir(original_cwd)
self.assertEqual(result, target_file)
def test_returns_none_when_file_not_found(self):
"""Given a filename that doesn't exist within max_depth, returns None."""
from comfyui_to_python.node_runtime import _find_file
original_cwd = os.getcwd()
with tempfile.TemporaryDirectory() as tmpdir:
try:
os.chdir(tmpdir)
result = _find_file("nonexistent_config.yaml")
finally:
os.chdir(original_cwd)
self.assertIsNone(result)
def test_returns_none_for_directory_not_file(self):
"""Given a directory with matching name, returns None (wants a file)."""
from comfyui_to_python.node_runtime import _find_file
original_cwd = os.getcwd()
with tempfile.TemporaryDirectory() as tmpdir:
# Create a directory (not file) with the target name
fake_dir = os.path.join(tmpdir, "extra_model_paths.yaml")
os.makedirs(fake_dir)
try:
os.chdir(tmpdir)
result = _find_file("extra_model_paths.yaml")
finally:
os.chdir(original_cwd)
self.assertIsNone(result)
class TestFindPathCwdFirst(unittest.TestCase):
"""Tests for find_path() checking CWD before walking to parent."""
def test_find_path_returns_cwd_when_cwd_is_match(self):
"""Given CWD is named 'ComfyUI', find_path('ComfyUI') returns CWD itself."""
from comfyui_to_python.node_runtime import find_path
original_cwd = os.getcwd()
with tempfile.TemporaryDirectory() as tmpdir:
comfyui_dir = os.path.join(tmpdir, "ComfyUI")
os.makedirs(comfyui_dir)
try:
os.chdir(comfyui_dir)
result = find_path("ComfyUI")
finally:
os.chdir(original_cwd)
self.assertEqual(result, comfyui_dir)
def test_find_path_checks_cwd_before_parent(self):
"""When CWD matches but parent also matches, CWD is returned (check-first)."""
from comfyui_to_python.node_runtime import find_path
original_cwd = os.getcwd()
with tempfile.TemporaryDirectory() as tmpdir:
# Create nested dirs: outer/ComfyUI/inner where inner is named ComfyUI too
outer_comfy = os.path.join(tmpdir, "ComfyUI")
os.makedirs(outer_comfy)
# Rename so we can create an inner dir also named ComfyUI
os.rename(outer_comfy, os.path.join(tmpdir, "ParentComfy"))
parent_comfy = os.path.join(tmpdir, "ParentComfy")
child_dir = os.path.join(parent_comfy, "subdir")
os.makedirs(child_dir)
try:
os.chdir(child_dir)
# When searching for ParentComfy, should find it at parent level
result = find_path("ParentComfy")
finally:
os.chdir(original_cwd)
self.assertEqual(result, parent_comfy)
class TestGetComfyuiPathThirdStrategy(unittest.TestCase):
"""Tests for get_comfyui_path() third-strategy verification."""
def test_fallback_rejects_non_comfyui_directory(self):
"""Given a directory named ComfyUI without nodes.py, third strategy rejects it."""
from comfyui_to_python.node_runtime import get_comfyui_path
with tempfile.TemporaryDirectory() as tmpdir:
# Create a fake "ComfyUI" directory (no nodes.py)
fake_dir = os.path.join(tmpdir, "ComfyUI")
os.makedirs(fake_dir)
original_cwd = os.getcwd()
try:
os.chdir(tmpdir)
with patch.dict("os.environ", {}, clear=False):
# Strip COMFYUI_PATH so fallback strategies activate
env_copy = dict(os.environ)
env_copy.pop("COMFYUI_PATH", None)
with patch.dict("os.environ", env_copy, clear=True):
result = get_comfyui_path()
finally:
os.chdir(original_cwd)
# Must NOT return the fake directory
self.assertNotEqual(result, fake_dir)
@unittest.skipUnless(os.path.isdir(_REAL_COMFYUI), "requires ../ComfyUI checkout")
class TestImportCustomNodesSysPathRestoration(unittest.TestCase):
"""Tests for sys.path restoration in import_custom_nodes()."""
def test_sys_path_restored_on_crash(self):
"""Given a crash during import_custom_nodes(), sys.path is still restored.
The crash fires on the "server" module load (after comfy/ filtering)
so that the try/finally sys.path restoration path is actually exercised.
"""
from comfyui_to_python.node_runtime import import_custom_nodes
original_path = sys.path[:]
call_order = []
def fake_load(name, path):
call_order.append(name)
if name == "server":
raise RuntimeError("simulated crash during server load")
return None # Let execution/nodes return None (falls through gracefully)
with (
patch(
"comfyui_to_python.node_runtime._load_module",
side_effect=fake_load,
),
patch(
"comfyui_to_python.node_runtime.get_comfyui_path",
return_value=_REAL_COMFYUI if os.path.isdir(_REAL_COMFYUI) else None,
),
):
try:
import_custom_nodes()
except RuntimeError:
pass # Expected crash
# When ComfyUI path is found, "execution" and "nodes" should be called
# before reaching the server module load inside try/finally.
self.assertIn("execution", call_order, "Expected 'execution' to be loaded")
self.assertIn("nodes", call_order, "Expected 'nodes' to be loaded")
self.assertIn("server", call_order, "Expected 'server' to trigger the crash")
# sys.path should still contain all original entries (no dangling state)
for entry in original_path:
self.assertIn(entry, sys.path)
if __name__ == "__main__":
unittest.main()
+16 -8
View File
@@ -146,9 +146,12 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
self.assertIn("def bootstrap_comfyui_runtime()", generated)
self.assertIn("def cleanup_comfyui_runtime(", generated)
self.assertIn("import comfy.options", generated)
self.assertIn("comfy.options.enable_args_parsing()", generated)
self.assertIn("import cuda_malloc", generated)
# Hardened imports use _load_module_temp() with temporary names,
# not bare 'import comfy.options' — avoids sys.modules pollution
self.assertIn("_load_module_temp", generated)
self.assertIn("options.py", generated)
self.assertIn("enable_args_parsing()", generated)
self.assertIn("cuda_malloc.py", generated)
self.assertNotIn("\nbootstrap_comfyui_runtime()\n", generated)
self.assertIn(
"def main(unload_models: bool | None = None):\n"
@@ -177,7 +180,9 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
main_section.index("add_extra_model_paths()"),
main_section.index("import torch"),
)
self.assertLess(generated.index("import cuda_malloc"), generated.index("import torch"))
# cuda_malloc is loaded via _load_module() inside bootstrap_comfyui_runtime()
# which runs before 'import torch' in main() — ordering preserved
self.assertIn("cuda_malloc", generated)
self.assertIn(
" finally:\n cleanup_comfyui_runtime(unload_models=unload_models)",
main_section,
@@ -383,10 +388,13 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
input_file.write_text(json.dumps(workflow), encoding="utf-8")
with patch(
"comfyui_to_python.get_node_class_mappings",
return_value={"LoadImage": LoadImage},
), patch("comfyui_to_python.import_custom_nodes"):
with (
patch(
"comfyui_to_python.get_node_class_mappings",
return_value={"LoadImage": LoadImage},
),
patch("comfyui_to_python.import_custom_nodes"),
):
run(
input_file=str(input_file),
output_file=str(output_file),
Generated
+35 -1
View File
@@ -63,8 +63,17 @@ dependencies = [
{ name = "black" },
]
[package.optional-dependencies]
dev = [
{ name = "ruff" },
]
[package.metadata]
requires-dist = [{ name = "black" }]
requires-dist = [
{ name = "black" },
{ name = "ruff", marker = "extra == 'dev'" },
]
provides-extras = ["dev"]
[[package]]
name = "mypy-extensions"
@@ -130,3 +139,28 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/39/2a/cbbf9250020a4a8dd53ba83a46c097b69e5eb49dd14e708f496f548c6612/pytokens-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:584c80c24b078eec1e227079d56dc22ff755e0ba8654d8383b2c549107528918", size = 116287, upload-time = "2026-01-30T01:03:30.912Z" },
{ url = "https://files.pythonhosted.org/packages/c6/78/397db326746f0a342855b81216ae1f0a32965deccfd7c830a2dbc66d2483/pytokens-0.4.1-py3-none-any.whl", hash = "sha256:26cef14744a8385f35d0e095dc8b3a7583f6c953c2e3d269c7f82484bf5ad2de", size = 13729, upload-time = "2026-01-30T01:03:45.029Z" },
]
[[package]]
name = "ruff"
version = "0.15.12"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/99/43/3291f1cc9106f4c63bdce7a8d0df5047fe8422a75b091c16b5e9355e0b11/ruff-0.15.12.tar.gz", hash = "sha256:ecea26adb26b4232c0c2ca19ccbc0083a68344180bba2a600605538ce51a40a6", size = 4643852, upload-time = "2026-04-24T18:17:14.305Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/c3/6e/e78ffb61d4686f3d96ba3df2c801161843746dcbcbb17a1e927d4829312b/ruff-0.15.12-py3-none-linux_armv6l.whl", hash = "sha256:f86f176e188e94d6bdbc09f09bfd9dc729059ad93d0e7390b5a73efe19f8861c", size = 10640713, upload-time = "2026-04-24T18:17:22.841Z" },
{ url = "https://files.pythonhosted.org/packages/ae/08/a317bc231fb9e7b93e4ef3089501e51922ff88d6936ce5cf870c4fe55419/ruff-0.15.12-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:e3bcd123364c3770b8e1b7baaf343cc99a35f197c5c6e8af79015c666c423a6c", size = 11069267, upload-time = "2026-04-24T18:17:30.105Z" },
{ url = "https://files.pythonhosted.org/packages/aa/a4/f828e9718d3dce1f5f11c39c4f65afd32783c8b2aebb2e3d259e492c47bd/ruff-0.15.12-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fe87510d000220aa1ed530d4448a7c696a0cae1213e5ec30e5874287b66557b5", size = 10397182, upload-time = "2026-04-24T18:17:07.177Z" },
{ url = "https://files.pythonhosted.org/packages/71/e0/3310fc6d1b5e1fdea22bf3b1b807c7e187b581021b0d7d4514cccdb5fb71/ruff-0.15.12-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:84a1630093121375a3e2a95b4a6dc7b59e2b4ee76216e32d81aae550a832d002", size = 10758012, upload-time = "2026-04-24T18:16:55.759Z" },
{ url = "https://files.pythonhosted.org/packages/11/c1/a606911aee04c324ddaa883ae418f3569792fd3c4a10c50e0dd0a2311e1e/ruff-0.15.12-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:fb129f40f114f089ebe0ca56c0d251cf2061b17651d464bb6478dc01e69f11f5", size = 10447479, upload-time = "2026-04-24T18:16:51.677Z" },
{ url = "https://files.pythonhosted.org/packages/9d/68/4201e8444f0894f21ab4aeeaee68aa4f10b51613514a20d80bd628d57e88/ruff-0.15.12-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b0c862b172d695db7598426b8af465e7e9ac00a3ea2a3630ee67eb82e366aaa6", size = 11234040, upload-time = "2026-04-24T18:17:16.529Z" },
{ url = "https://files.pythonhosted.org/packages/34/ff/8a6d6cf4ccc23fd67060874e832c18919d1557a0611ebef03fdb01fff11e/ruff-0.15.12-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2849ea9f3484c3aca43a82f484210370319e7170df4dfe4843395ddf6c57bc33", size = 12087377, upload-time = "2026-04-24T18:17:04.944Z" },
{ url = "https://files.pythonhosted.org/packages/85/f6/c669cf73f5152f623d34e69866a46d5e6185816b19fcd5b6dd8a2d299922/ruff-0.15.12-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9e77c7e51c07fe396826d5969a5b846d9cd4c402535835fb6e21ce8b28fef847", size = 11367784, upload-time = "2026-04-24T18:17:25.409Z" },
{ url = "https://files.pythonhosted.org/packages/e8/39/c61d193b8a1daaa8977f7dea9e8d8ba866e02ea7b65d32f6861693aa4c12/ruff-0.15.12-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:83b2f4f2f3b1026b5fb449b467d9264bf22067b600f7b6f41fc5958909f449d0", size = 11344088, upload-time = "2026-04-24T18:17:12.258Z" },
{ url = "https://files.pythonhosted.org/packages/c2/8d/49afab3645e31e12c590acb6d3b5b69d7aab5b81926dbaf7461f9441f37a/ruff-0.15.12-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:9ba3b8f1afd7e2e43d8943e55f249e13f9682fde09711644a6e7290eb4f3e339", size = 11271770, upload-time = "2026-04-24T18:17:02.457Z" },
{ url = "https://files.pythonhosted.org/packages/46/06/33f41fe94403e2b755481cdfb9b7ef3e4e0ed031c4581124658d935d52b4/ruff-0.15.12-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:e852ba9fdc890655e1d78f2df1499efbe0e54126bd405362154a75e2bde159c5", size = 10719355, upload-time = "2026-04-24T18:17:27.648Z" },
{ url = "https://files.pythonhosted.org/packages/0d/59/18aa4e014debbf559670e4048e39260a85c7fcee84acfd761ac01e7b8d35/ruff-0.15.12-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:dd8aed930da53780d22fc70bdf84452c843cf64f8cb4eb38984319c24c5cd5fd", size = 10462758, upload-time = "2026-04-24T18:17:32.347Z" },
{ url = "https://files.pythonhosted.org/packages/25/e7/cc9f16fd0f3b5fddcbd7ec3d6ae30c8f3fde1047f32a4093a98d633c6570/ruff-0.15.12-py3-none-musllinux_1_2_i686.whl", hash = "sha256:01da3988d225628b709493d7dc67c3b9b12c0210016b08690ef9bd27970b262b", size = 10953498, upload-time = "2026-04-24T18:17:20.674Z" },
{ url = "https://files.pythonhosted.org/packages/72/7a/a9ba7f98c7a575978698f4230c5e8cc54bbc761af34f560818f933dafa0c/ruff-0.15.12-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:9cae0f92bd5700d1213188b31cd3bdd2b315361296d10b96b8e2337d3d11f53e", size = 11447765, upload-time = "2026-04-24T18:17:09.755Z" },
{ url = "https://files.pythonhosted.org/packages/ea/f9/0ae446942c846b8266059ad8a30702a35afae55f5cdc54c5adf8d7afdc27/ruff-0.15.12-py3-none-win32.whl", hash = "sha256:d0185894e038d7043ba8fd6aee7499ece6462dc0ea9f1e260c7451807c714c20", size = 10657277, upload-time = "2026-04-24T18:17:18.591Z" },
{ url = "https://files.pythonhosted.org/packages/33/f1/9614e03e1cdcbf9437570b5400ced8a720b5db22b28d8e0f1bda429f660d/ruff-0.15.12-py3-none-win_amd64.whl", hash = "sha256:c87a162d61ab3adca47c03f7f717c68672edec7d1b5499e652331780fe74950d", size = 11837758, upload-time = "2026-04-24T18:17:00.113Z" },
{ url = "https://files.pythonhosted.org/packages/c0/98/6beb4b351e472e5f4c4613f7c35a5290b8be2497e183825310c4c3a3984b/ruff-0.15.12-py3-none-win_arm64.whl", hash = "sha256:a538f7a82d061cee7be55542aca1d86d1393d55d81d4fcc314370f4340930d4f", size = 11120821, upload-time = "2026-04-24T18:16:57.979Z" },
]