74 lines
2.5 KiB
Python
74 lines
2.5 KiB
Python
import os
|
|
import sys
|
|
import contextlib
|
|
import importlib
|
|
import traceback
|
|
import folder_paths
|
|
|
|
@contextlib.contextmanager
|
|
def suppress_console():
|
|
with open(os.devnull, 'w') as devnull:
|
|
old_stdout = sys.stdout
|
|
old_stderr = sys.stderr
|
|
sys.stdout = devnull
|
|
sys.stderr = devnull
|
|
try:
|
|
yield
|
|
finally:
|
|
sys.stdout = old_stdout
|
|
sys.stderr = old_stderr
|
|
|
|
def get_custom_node_mappings(module_path, ignore=set()):
|
|
mappings = []
|
|
|
|
module_name = os.path.basename(module_path)
|
|
if os.path.isfile(module_path):
|
|
sp = os.path.splitext(module_path)
|
|
module_name = sp[0]
|
|
try:
|
|
if os.path.isfile(module_path):
|
|
module_spec = importlib.util.spec_from_file_location(module_name, module_path)
|
|
module_dir = os.path.split(module_path)[0]
|
|
else:
|
|
module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py"))
|
|
module_dir = module_path
|
|
|
|
module = importlib.util.module_from_spec(module_spec)
|
|
module_spec.loader.exec_module(module)
|
|
|
|
if hasattr(module, "NODE_CLASS_MAPPINGS") and getattr(module, "NODE_CLASS_MAPPINGS") is not None:
|
|
for name in module.NODE_CLASS_MAPPINGS:
|
|
if name not in ignore:
|
|
mappings.append(name)
|
|
|
|
return (module_name, mappings)
|
|
else:
|
|
# Skip module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.
|
|
return (module_name, [])
|
|
except Exception as e:
|
|
return (module_name, [])
|
|
|
|
def get_custom_node_list():
|
|
"""
|
|
Returns a list of custom nodes.
|
|
"""
|
|
custom_nodes = {}
|
|
node_paths = folder_paths.get_folder_paths("custom_nodes")
|
|
for custom_node_path in node_paths:
|
|
possible_modules = os.listdir(os.path.realpath(custom_node_path))
|
|
if "__pycache__" in possible_modules:
|
|
possible_modules.remove("__pycache__")
|
|
|
|
for possible_module in possible_modules:
|
|
module_path = os.path.join(custom_node_path, possible_module)
|
|
if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue
|
|
if module_path.endswith(".disabled"): continue
|
|
mappings = get_custom_node_mappings(module_path)
|
|
custom_nodes[mappings[0]] = mappings[1]
|
|
|
|
return custom_nodes
|
|
|
|
def get_custom_node_list_silent():
|
|
with suppress_console():
|
|
return get_custom_node_list()
|