|
| 1 | +import os |
| 2 | +from pathlib import Path |
| 3 | +from typing import Optional |
| 4 | + |
| 5 | +from pydantic_settings import PydanticBaseSettingsSource, TomlConfigSettingsSource |
| 6 | + |
| 7 | +from comfy_config.types import ( |
| 8 | + ComfyConfig, |
| 9 | + ProjectConfig, |
| 10 | + PyProjectConfig, |
| 11 | + PyProjectSettings |
| 12 | +) |
| 13 | + |
| 14 | +""" |
| 15 | +Extract configuration from a custom node directory's pyproject.toml file or a Python file. |
| 16 | +
|
| 17 | +This function reads and parses the pyproject.toml file in the specified directory |
| 18 | +to extract project and ComfyUI-specific configuration information. If no |
| 19 | +pyproject.toml file is found, it creates a minimal configuration using the |
| 20 | +folder name as the project name. If a Python file is provided, it uses the |
| 21 | +file name (without extension) as the project name. |
| 22 | +
|
| 23 | +Args: |
| 24 | + path (str): Path to the directory containing the pyproject.toml file, or |
| 25 | + path to a .py file. If pyproject.toml doesn't exist in a directory, |
| 26 | + the folder name will be used as the default project name. If a .py |
| 27 | + file is provided, the filename (without .py extension) will be used |
| 28 | + as the project name. |
| 29 | +
|
| 30 | +Returns: |
| 31 | + Optional[PyProjectConfig]: A PyProjectConfig object containing: |
| 32 | + - project: Basic project information (name, version, dependencies, etc.) |
| 33 | + - tool_comfy: ComfyUI-specific configuration (publisher_id, models, etc.) |
| 34 | + Returns None if configuration extraction fails or if the provided file |
| 35 | + is not a Python file. |
| 36 | +
|
| 37 | +Notes: |
| 38 | + - If pyproject.toml is missing in a directory, creates a default config with folder name |
| 39 | + - If a .py file is provided, creates a default config with filename (without extension) |
| 40 | + - Returns None for non-Python files |
| 41 | +
|
| 42 | +Example: |
| 43 | + >>> from comfy_config import config_parser |
| 44 | + >>> # For directory |
| 45 | + >>> custom_node_dir = os.path.dirname(os.path.realpath(__file__)) |
| 46 | + >>> project_config = config_parser.extract_node_configuration(custom_node_dir) |
| 47 | + >>> print(project_config.project.name) # "my_custom_node" or name from pyproject.toml |
| 48 | + >>> |
| 49 | + >>> # For single-file Python node file |
| 50 | + >>> py_file_path = os.path.realpath(__file__) # "/path/to/my_node.py" |
| 51 | + >>> project_config = config_parser.extract_node_configuration(py_file_path) |
| 52 | + >>> print(project_config.project.name) # "my_node" |
| 53 | +""" |
| 54 | +def extract_node_configuration(path) -> Optional[PyProjectConfig]: |
| 55 | + if os.path.isfile(path): |
| 56 | + file_path = Path(path) |
| 57 | + |
| 58 | + if file_path.suffix.lower() != '.py': |
| 59 | + return None |
| 60 | + |
| 61 | + project_name = file_path.stem |
| 62 | + project = ProjectConfig(name=project_name) |
| 63 | + comfy = ComfyConfig() |
| 64 | + return PyProjectConfig(project=project, tool_comfy=comfy) |
| 65 | + |
| 66 | + folder_name = os.path.basename(path) |
| 67 | + toml_path = Path(path) / "pyproject.toml" |
| 68 | + |
| 69 | + if not toml_path.exists(): |
| 70 | + project = ProjectConfig(name=folder_name) |
| 71 | + comfy = ComfyConfig() |
| 72 | + return PyProjectConfig(project=project, tool_comfy=comfy) |
| 73 | + |
| 74 | + raw_settings = load_pyproject_settings(toml_path) |
| 75 | + |
| 76 | + project_data = raw_settings.project |
| 77 | + |
| 78 | + tool_data = raw_settings.tool |
| 79 | + comfy_data = tool_data.get("comfy", {}) if tool_data else {} |
| 80 | + |
| 81 | + return PyProjectConfig(project=project_data, tool_comfy=comfy_data) |
| 82 | + |
| 83 | + |
| 84 | +def load_pyproject_settings(toml_path: Path) -> PyProjectSettings: |
| 85 | + class PyProjectLoader(PyProjectSettings): |
| 86 | + @classmethod |
| 87 | + def settings_customise_sources( |
| 88 | + cls, |
| 89 | + settings_cls, |
| 90 | + init_settings: PydanticBaseSettingsSource, |
| 91 | + env_settings: PydanticBaseSettingsSource, |
| 92 | + dotenv_settings: PydanticBaseSettingsSource, |
| 93 | + file_secret_settings: PydanticBaseSettingsSource, |
| 94 | + ): |
| 95 | + return (TomlConfigSettingsSource(settings_cls, toml_path),) |
| 96 | + |
| 97 | + return PyProjectLoader() |
0 commit comments