diff --git a/comfy_config/config_parser.py b/comfy_config/config_parser.py index dfea58172..5c59ab4f2 100644 --- a/comfy_config/config_parser.py +++ b/comfy_config/config_parser.py @@ -1,23 +1,17 @@ import os from typing import Optional -import tomlkit -import tomlkit.exceptions +import tomllib +from pydantic import ValidationError import logging from comfy_config.types import ( - ComfyConfig, - License, - Model, ProjectConfig, PyProjectConfig, - URLs, ) """ -Original implementation comes from https://github.com/Comfy-Org/comfy-cli/blob/2e36f33dd39ef43b5acf7d1fc5acc5e01be92360/comfy_cli/registry/config_parser.py#L146 - Extract configuration from a custom node directory's pyproject.toml file. This function reads and parses the pyproject.toml file in the specified directory @@ -46,70 +40,42 @@ Example: >>> print(project_config.project.name) # "my_custom_node" or name from pyproject.toml >>> nodes.EXTENSION_WEB_DIRS[project_config.project.name] = js_dir """ -def extract_node_configuration( - path, -) -> Optional[PyProjectConfig]: + + +def extract_node_configuration(path) -> Optional[PyProjectConfig]: folder_name = os.path.basename(path) + toml_path = os.path.join(path, "pyproject.toml") - path = os.path.join(path, "pyproject.toml") - - if not os.path.isfile(path): - logging.warning("No pyproject.toml file found in the current directory, will use custom node folder name as project name as default.") - - project = ProjectConfig( - name=folder_name, - ) - - return PyProjectConfig(project=project) - - with open(path, "r") as file: - data = tomlkit.load(file) - - project_data = data.get("project", {}) - urls_data = project_data.get("urls", {}) - comfy_data = data.get("tool", {}).get("comfy", {}) - - license_data = project_data.get("license", {}) - if isinstance(license_data, str): - license = License(text=license_data) + if not os.path.isfile(toml_path): logging.warning( - 'Warning: License should be in one of these two formats: license = {file = "LICENSE"} OR license = {text = "MIT License"}. Please check the documentation: https://docs.comfy.org/registry/specifications.' - ) - elif isinstance(license_data, dict): - if "file" in license_data or "text" in license_data: - license = License(file=license_data.get("file", ""), text=license_data.get("text", "")) - else: - logging.warning( - 'Warning: License should be in one of these two formats: license = {file = "LICENSE"} OR license = {text = "MIT License"}. Please check the documentation: https://docs.comfy.org/registry/specifications.' - ) - license = License() - else: - license = License() - logging.warning( - 'Warning: License should be in one of these two formats: license = {file = "LICENSE"} OR license = {text = "MIT License"}. Please check the documentation: https://docs.comfy.org/registry/specifications.' - ) + "No pyproject.toml file found in the current directory, will use custom node folder name as project name as default.") - project = ProjectConfig( - name=project_data.get("name", ""), - description=project_data.get("description", ""), - version=project_data.get("version", ""), - requires_python=project_data.get("requires-python", ""), - dependencies=project_data.get("dependencies", []), - license=license, - urls=URLs( - homepage=urls_data.get("Homepage", ""), - documentation=urls_data.get("Documentation", ""), - repository=urls_data.get("Repository", ""), - issues=urls_data.get("Issues", ""), - ), - ) + try: + project = ProjectConfig(name=folder_name) + return PyProjectConfig(project=project) + except ValidationError as e: + logging.error(f"Failed to create default configuration: {e}") + return None - comfy = ComfyConfig( - publisher_id=comfy_data.get("PublisherId", ""), - display_name=comfy_data.get("DisplayName", ""), - icon=comfy_data.get("Icon", ""), - models=[Model(location=m["location"], model_url=m["model_url"]) for m in comfy_data.get("Models", [])], - includes=comfy_data.get("includes", []), - ) + try: + with open(toml_path, "rb") as f: + data = tomllib.load(f) + except Exception as e: + logging.error(f"Failed to read pyproject.toml: {e}") + return None - return PyProjectConfig(project=project, tool_comfy=comfy) + try: + config_data = { + "project": data.get("project", {}), + "tool_comfy": data.get("tool", {}).get("comfy", {}) + } + + return PyProjectConfig(**config_data) + + except ValidationError as e: + logging.error(f"Validation error while parsing configuration: {e}") + logging.error(f"Validation details: {e.errors()}") + return None + except Exception as e: + logging.error(f"Unexpected error while parsing configuration: {e}") + return None diff --git a/comfy_config/types.py b/comfy_config/types.py index f56b5f4c1..8657ac13c 100644 --- a/comfy_config/types.py +++ b/comfy_config/types.py @@ -1,12 +1,11 @@ -from dataclasses import dataclass, field +from pydantic import BaseModel, Field from typing import List, Optional # IMPORTANT: The type definitions specified in pyproject.toml for custom nodes # must remain synchronized with the corresponding files in the https://github.com/Comfy-Org/comfy-cli/blob/main/comfy_cli/registry/types.py. # Any changes to one must be reflected in the other to maintain consistency. -@dataclass -class NodeVersion: +class NodeVersion(BaseModel): changelog: str dependencies: List[str] deprecated: bool @@ -15,8 +14,7 @@ class NodeVersion: download_url: str -@dataclass -class Node: +class Node(BaseModel): id: str name: str description: str @@ -24,57 +22,50 @@ class Node: license: Optional[str] = None icon: Optional[str] = None repository: Optional[str] = None - tags: List[str] = field(default_factory=list) + tags: List[str] = Field(default_factory=list) latest_version: Optional[NodeVersion] = None -@dataclass -class PublishNodeVersionResponse: +class PublishNodeVersionResponse(BaseModel): node_version: NodeVersion signedUrl: str -@dataclass -class URLs: - homepage: str = "" - documentation: str = "" - repository: str = "" - issues: str = "" +class URLs(BaseModel): + homepage: str = Field(default="", alias="Homepage") + documentation: str = Field(default="", alias="Documentation") + repository: str = Field(default="", alias="Repository") + issues: str = Field(default="", alias="Issues") -@dataclass -class Model: +class Model(BaseModel): location: str model_url: str -@dataclass -class ComfyConfig: - publisher_id: str = "" - display_name: str = "" - icon: str = "" - models: List[Model] = field(default_factory=list) - includes: List[str] = field(default_factory=list) +class ComfyConfig(BaseModel): + publisher_id: str = Field(default="", alias="PublisherId") + display_name: str = Field(default="", alias="DisplayName") + icon: str = Field(default="", alias="Icon") + models: List[Model] = Field(default_factory=list, alias="Models") + includes: List[str] = Field(default_factory=list) -@dataclass -class License: +class License(BaseModel): file: str = "" text: str = "" -@dataclass -class ProjectConfig: +class ProjectConfig(BaseModel): name: str = "" description: str = "" version: str = "1.0.0" - requires_python: str = ">= 3.9" - dependencies: List[str] = field(default_factory=list) - license: License = field(default_factory=License) - urls: URLs = field(default_factory=URLs) + requires_python: str = Field(default=">= 3.9", alias="requires-python") + dependencies: List[str] = Field(default_factory=list) + license: License = Field(default_factory=License) + urls: URLs = Field(default_factory=URLs) -@dataclass -class PyProjectConfig: - project: ProjectConfig = field(default_factory=ProjectConfig) - tool_comfy: ComfyConfig = field(default_factory=ComfyConfig) +class PyProjectConfig(BaseModel): + project: ProjectConfig = Field(default_factory=ProjectConfig) + tool_comfy: ComfyConfig = Field(default_factory=ComfyConfig) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index f51c80e54..3e2991563 100644 --- a/requirements.txt +++ b/requirements.txt @@ -24,4 +24,3 @@ spandrel soundfile av>=14.2.0 pydantic~=2.0 -tomlkit