Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 110 additions & 10 deletions app/custom_node_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,93 @@
import folder_paths
import glob
from aiohttp import web
import json
import logging
from functools import lru_cache

from utils.json_util import merge_json_recursive


# Extra locale files to load into main.json
EXTRA_LOCALE_FILES = [
"nodeDefs.json",
"commands.json",
"settings.json",
]


def safe_load_json_file(file_path: str) -> dict:
if not os.path.exists(file_path):
return {}

try:
with open(file_path, "r", encoding="utf-8") as f:
return json.load(f)
except json.JSONDecodeError:
logging.error(f"Error loading {file_path}")
return {}


class CustomNodeManager:
"""
Placeholder to refactor the custom node management features from ComfyUI-Manager.
Currently it only contains the custom workflow templates feature.
"""
@lru_cache(maxsize=1)
def build_translations(self):
"""Load all custom nodes translations during initialization. Translations are
expected to be loaded from `locales/` folder.

The folder structure is expected to be the following:
- custom_nodes/
- custom_node_1/
- locales/
- en/
- main.json
- commands.json
- settings.json

returned translations are expected to be in the following format:
{
"en": {
"nodeDefs": {...},
"commands": {...},
"settings": {...},
...{other main.json keys}
}
}
"""

translations = {}

for folder in folder_paths.get_folder_paths("custom_nodes"):
# Sort glob results for deterministic ordering
for custom_node_dir in sorted(glob.glob(os.path.join(folder, "*/"))):
locales_dir = os.path.join(custom_node_dir, "locales")
if not os.path.exists(locales_dir):
continue

for lang_dir in glob.glob(os.path.join(locales_dir, "*/")):
lang_code = os.path.basename(os.path.dirname(lang_dir))

if lang_code not in translations:
translations[lang_code] = {}

# Load main.json
main_file = os.path.join(lang_dir, "main.json")
node_translations = safe_load_json_file(main_file)

# Load extra locale files
for extra_file in EXTRA_LOCALE_FILES:
extra_file_path = os.path.join(lang_dir, extra_file)
key = extra_file.split(".")[0]
json_data = safe_load_json_file(extra_file_path)
if json_data:
node_translations[key] = json_data

if node_translations:
translations[lang_code] = merge_json_recursive(
translations[lang_code], node_translations
)

return translations

def add_routes(self, routes, webapp, loadedModules):

@routes.get("/workflow_templates")
Expand All @@ -18,17 +99,36 @@ async def get_workflow_templates(request):
files = [
file
for folder in folder_paths.get_folder_paths("custom_nodes")
for file in glob.glob(os.path.join(folder, '*/example_workflows/*.json'))
for file in glob.glob(
os.path.join(folder, "*/example_workflows/*.json")
)
]
workflow_templates_dict = {} # custom_nodes folder name -> example workflow names
workflow_templates_dict = (
{}
) # custom_nodes folder name -> example workflow names
for file in files:
custom_nodes_name = os.path.basename(os.path.dirname(os.path.dirname(file)))
custom_nodes_name = os.path.basename(
os.path.dirname(os.path.dirname(file))
)
workflow_name = os.path.splitext(os.path.basename(file))[0]
workflow_templates_dict.setdefault(custom_nodes_name, []).append(workflow_name)
workflow_templates_dict.setdefault(custom_nodes_name, []).append(
workflow_name
)
return web.json_response(workflow_templates_dict)

# Serve workflow templates from custom nodes.
for module_name, module_dir in loadedModules:
workflows_dir = os.path.join(module_dir, 'example_workflows')
workflows_dir = os.path.join(module_dir, "example_workflows")
if os.path.exists(workflows_dir):
webapp.add_routes([web.static('/api/workflow_templates/' + module_name, workflows_dir)])
webapp.add_routes(
[
web.static(
"/api/workflow_templates/" + module_name, workflows_dir
)
]
)

@routes.get("/i18n")
async def get_i18n(request):
"""Returns translations from all custom nodes' locales folders."""
return web.json_response(self.build_translations())
121 changes: 114 additions & 7 deletions tests-unit/app_test/custom_node_manager_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,39 +2,146 @@
from aiohttp import web
from unittest.mock import patch
from app.custom_node_manager import CustomNodeManager
import json

pytestmark = (
pytest.mark.asyncio
) # This applies the asyncio mark to all test functions in the module


@pytest.fixture
def custom_node_manager():
return CustomNodeManager()


@pytest.fixture
def app(custom_node_manager):
app = web.Application()
routes = web.RouteTableDef()
custom_node_manager.add_routes(routes, app, [("ComfyUI-TestExtension1", "ComfyUI-TestExtension1")])
custom_node_manager.add_routes(
routes, app, [("ComfyUI-TestExtension1", "ComfyUI-TestExtension1")]
)
app.add_routes(routes)
return app


async def test_get_workflow_templates(aiohttp_client, app, tmp_path):
client = await aiohttp_client(app)
# Setup temporary custom nodes file structure with 1 workflow file
custom_nodes_dir = tmp_path / "custom_nodes"
example_workflows_dir = custom_nodes_dir / "ComfyUI-TestExtension1" / "example_workflows"
example_workflows_dir = (
custom_nodes_dir / "ComfyUI-TestExtension1" / "example_workflows"
)
example_workflows_dir.mkdir(parents=True)
template_file = example_workflows_dir / "workflow1.json"
template_file.write_text('')
template_file.write_text("")

with patch('folder_paths.folder_names_and_paths', {
'custom_nodes': ([str(custom_nodes_dir)], None)
}):
response = await client.get('/workflow_templates')
with patch(
"folder_paths.folder_names_and_paths",
{"custom_nodes": ([str(custom_nodes_dir)], None)},
):
response = await client.get("/workflow_templates")
assert response.status == 200
workflows_dict = await response.json()
assert isinstance(workflows_dict, dict)
assert "ComfyUI-TestExtension1" in workflows_dict
assert isinstance(workflows_dict["ComfyUI-TestExtension1"], list)
assert workflows_dict["ComfyUI-TestExtension1"][0] == "workflow1"


async def test_build_translations_empty_when_no_locales(custom_node_manager, tmp_path):
custom_nodes_dir = tmp_path / "custom_nodes"
custom_nodes_dir.mkdir(parents=True)

with patch("folder_paths.get_folder_paths", return_value=[str(custom_nodes_dir)]):
translations = custom_node_manager.build_translations()
assert translations == {}


async def test_build_translations_loads_all_files(custom_node_manager, tmp_path):
# Setup test directory structure
custom_nodes_dir = tmp_path / "custom_nodes" / "test-extension"
locales_dir = custom_nodes_dir / "locales" / "en"
locales_dir.mkdir(parents=True)

# Create test translation files
main_content = {"title": "Test Extension"}
(locales_dir / "main.json").write_text(json.dumps(main_content))

node_defs = {"node1": "Node 1"}
(locales_dir / "nodeDefs.json").write_text(json.dumps(node_defs))

commands = {"cmd1": "Command 1"}
(locales_dir / "commands.json").write_text(json.dumps(commands))

settings = {"setting1": "Setting 1"}
(locales_dir / "settings.json").write_text(json.dumps(settings))

with patch(
"folder_paths.get_folder_paths", return_value=[tmp_path / "custom_nodes"]
):
translations = custom_node_manager.build_translations()

assert translations == {
"en": {
"title": "Test Extension",
"nodeDefs": {"node1": "Node 1"},
"commands": {"cmd1": "Command 1"},
"settings": {"setting1": "Setting 1"},
}
}


async def test_build_translations_handles_invalid_json(custom_node_manager, tmp_path):
# Setup test directory structure
custom_nodes_dir = tmp_path / "custom_nodes" / "test-extension"
locales_dir = custom_nodes_dir / "locales" / "en"
locales_dir.mkdir(parents=True)

# Create valid main.json
main_content = {"title": "Test Extension"}
(locales_dir / "main.json").write_text(json.dumps(main_content))

# Create invalid JSON file
(locales_dir / "nodeDefs.json").write_text("invalid json{")

with patch(
"folder_paths.get_folder_paths", return_value=[tmp_path / "custom_nodes"]
):
translations = custom_node_manager.build_translations()

assert translations == {
"en": {
"title": "Test Extension",
}
}


async def test_build_translations_merges_multiple_extensions(
custom_node_manager, tmp_path
):
# Setup test directory structure for two extensions
custom_nodes_dir = tmp_path / "custom_nodes"
ext1_dir = custom_nodes_dir / "extension1" / "locales" / "en"
ext2_dir = custom_nodes_dir / "extension2" / "locales" / "en"
ext1_dir.mkdir(parents=True)
ext2_dir.mkdir(parents=True)

# Create translation files for extension 1
ext1_main = {"title": "Extension 1", "shared": "Original"}
(ext1_dir / "main.json").write_text(json.dumps(ext1_main))

# Create translation files for extension 2
ext2_main = {"description": "Extension 2", "shared": "Override"}
(ext2_dir / "main.json").write_text(json.dumps(ext2_main))

with patch("folder_paths.get_folder_paths", return_value=[str(custom_nodes_dir)]):
translations = custom_node_manager.build_translations()

assert translations == {
"en": {
"title": "Extension 1",
"description": "Extension 2",
"shared": "Override", # Second extension should override first
}
}
71 changes: 71 additions & 0 deletions tests-unit/utils/json_util_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
from utils.json_util import merge_json_recursive


def test_merge_simple_dicts():
base = {"a": 1, "b": 2}
update = {"b": 3, "c": 4}
expected = {"a": 1, "b": 3, "c": 4}
assert merge_json_recursive(base, update) == expected


def test_merge_nested_dicts():
base = {"a": {"x": 1, "y": 2}, "b": 3}
update = {"a": {"y": 4, "z": 5}}
expected = {"a": {"x": 1, "y": 4, "z": 5}, "b": 3}
assert merge_json_recursive(base, update) == expected


def test_merge_lists():
base = {"a": [1, 2], "b": 3}
update = {"a": [3, 4]}
expected = {"a": [1, 2, 3, 4], "b": 3}
assert merge_json_recursive(base, update) == expected


def test_merge_nested_lists():
base = {"a": {"x": [1, 2]}}
update = {"a": {"x": [3, 4]}}
expected = {"a": {"x": [1, 2, 3, 4]}}
assert merge_json_recursive(base, update) == expected


def test_merge_mixed_types():
base = {"a": [1, 2], "b": {"x": 1}}
update = {"a": [3], "b": {"y": 2}}
expected = {"a": [1, 2, 3], "b": {"x": 1, "y": 2}}
assert merge_json_recursive(base, update) == expected


def test_merge_overwrite_non_dict():
base = {"a": 1}
update = {"a": {"x": 2}}
expected = {"a": {"x": 2}}
assert merge_json_recursive(base, update) == expected


def test_merge_empty_dicts():
base = {}
update = {"a": 1}
expected = {"a": 1}
assert merge_json_recursive(base, update) == expected


def test_merge_none_values():
base = {"a": None}
update = {"a": {"x": 1}}
expected = {"a": {"x": 1}}
assert merge_json_recursive(base, update) == expected


def test_merge_different_types():
base = {"a": [1, 2]}
update = {"a": "string"}
expected = {"a": "string"}
assert merge_json_recursive(base, update) == expected


def test_merge_complex_nested():
base = {"a": [1, 2], "b": {"x": [3, 4], "y": {"p": 1}}}
update = {"a": [5], "b": {"x": [6], "y": {"q": 2}}}
expected = {"a": [1, 2, 5], "b": {"x": [3, 4, 6], "y": {"p": 1, "q": 2}}}
assert merge_json_recursive(base, update) == expected
26 changes: 26 additions & 0 deletions utils/json_util.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
def merge_json_recursive(base, update):
"""Recursively merge two JSON-like objects.
- Dictionaries are merged recursively
- Lists are concatenated
- Other types are overwritten by the update value

Args:
base: Base JSON-like object
update: Update JSON-like object to merge into base

Returns:
Merged JSON-like object
"""
if not isinstance(base, dict) or not isinstance(update, dict):
if isinstance(base, list) and isinstance(update, list):
return base + update
return update

merged = base.copy()
for key, value in update.items():
if key in merged:
merged[key] = merge_json_recursive(merged[key], value)
else:
merged[key] = value

return merged
Loading