[i18n] Add /i18n endpoint to provide all custom node translations (#6558)

huchenlei · web-flow · commit a058f520901e · 2025-01-22T17:15:45.000-05:00
* [i18n] Add /i18n endpoint to provide all custom node translations

* Sort glob result for deterministic ordering

* Update comment
diff --git a/app/custom_node_manager.py b/app/custom_node_manager.py
@@ -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")
@@ -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())
diff --git a/tests-unit/app_test/custom_node_manager_test.py b/tests-unit/app_test/custom_node_manager_test.py
@@ -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
+            }
+        }
diff --git a/tests-unit/utils/json_util_test.py b/tests-unit/utils/json_util_test.py
@@ -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
diff --git a/utils/json_util.py b/utils/json_util.py
@@ -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
