use pydantic instead

comfy_config/config_parser.py

@@ -1,23 +1,17 @@
 import os
 from typing import Optional
 
-import tomlkit
+import tomllib
-import tomlkit.exceptions
 
+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(toml_path):
-
-    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)
         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.")
-        )
-    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.'
-        )
 
-    project = ProjectConfig(
+        try:
-        name=project_data.get("name", ""),
+            project = ProjectConfig(name=folder_name)
-        description=project_data.get("description", ""),
+            return PyProjectConfig(project=project)
-        version=project_data.get("version", ""),
+        except ValidationError as e:
-        requires_python=project_data.get("requires-python", ""),
+            logging.error(f"Failed to create default configuration: {e}")
-        dependencies=project_data.get("dependencies", []),
+            return None
-        license=license,
-        urls=URLs(
-            homepage=urls_data.get("Homepage", ""),
-            documentation=urls_data.get("Documentation", ""),
-            repository=urls_data.get("Repository", ""),
-            issues=urls_data.get("Issues", ""),
-        ),
-    )
 
-    comfy = ComfyConfig(
+    try:
-        publisher_id=comfy_data.get("PublisherId", ""),
+        with open(toml_path, "rb") as f:
-        display_name=comfy_data.get("DisplayName", ""),
+            data = tomllib.load(f)
-        icon=comfy_data.get("Icon", ""),
+    except Exception as e:
-        models=[Model(location=m["location"], model_url=m["model_url"]) for m in comfy_data.get("Models", [])],
+        logging.error(f"Failed to read pyproject.toml: {e}")
-        includes=comfy_data.get("includes", []),
+        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

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(BaseModel):
-class NodeVersion:
     changelog: str
     dependencies: List[str]
     deprecated: bool
@@ -15,8 +14,7 @@ class NodeVersion:
     download_url: str
 
 
-@dataclass
+class Node(BaseModel):
-class Node:
     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(BaseModel):
-class PublishNodeVersionResponse:
     node_version: NodeVersion
     signedUrl: str
 
 
-@dataclass
+class URLs(BaseModel):
-class URLs:
+    homepage: str = Field(default="", alias="Homepage")
-    homepage: str = ""
+    documentation: str = Field(default="", alias="Documentation")
-    documentation: str = ""
+    repository: str = Field(default="", alias="Repository")
-    repository: str = ""
+    issues: str = Field(default="", alias="Issues")
-    issues: str = ""
 
 
-@dataclass
+class Model(BaseModel):
-class Model:
     location: str
     model_url: str
 
 
-@dataclass
+class ComfyConfig(BaseModel):
-class ComfyConfig:
+    publisher_id: str = Field(default="", alias="PublisherId")
-    publisher_id: str = ""
+    display_name: str = Field(default="", alias="DisplayName")
-    display_name: str = ""
+    icon: str = Field(default="", alias="Icon")
-    icon: str = ""
+    models: List[Model] = Field(default_factory=list, alias="Models")
-    models: List[Model] = field(default_factory=list)
+    includes: List[str] = Field(default_factory=list)
-    includes: List[str] = field(default_factory=list)
 
 
-@dataclass
+class License(BaseModel):
-class License:
     file: str = ""
     text: str = ""
 
 
-@dataclass
+class ProjectConfig(BaseModel):
-class ProjectConfig:
     name: str = ""
     description: str = ""
     version: str = "1.0.0"
-    requires_python: str = ">= 3.9"
+    requires_python: str = Field(default=">= 3.9", alias="requires-python")
-    dependencies: List[str] = field(default_factory=list)
+    dependencies: List[str] = Field(default_factory=list)
-    license: License = field(default_factory=License)
+    license: License = Field(default_factory=License)
-    urls: URLs = field(default_factory=URLs)
+    urls: URLs = Field(default_factory=URLs)
 
 
-@dataclass
+class PyProjectConfig(BaseModel):
-class PyProjectConfig:
+    project: ProjectConfig = Field(default_factory=ProjectConfig)
-    project: ProjectConfig = field(default_factory=ProjectConfig)
+    tool_comfy: ComfyConfig = Field(default_factory=ComfyConfig)
-    tool_comfy: ComfyConfig = field(default_factory=ComfyConfig)

requirements.txt

@@ -24,4 +24,3 @@ spandrel
 soundfile
 av>=14.2.0
 pydantic~=2.0
-tomlkit