Refactor Registry API code (#40)

* Refactor Registry API code

* fmt

comfy_cli/command/custom_nodes/command.py

@@ -11,9 +11,8 @@
 from comfy_cli.config_manager import ConfigManager
 from comfy_cli.workspace_manager import WorkspaceManager
 
-
 from comfy_cli.registry import (
-    publish_node_version,
+    RegistryAPI,
     extract_node_configuration,
     upload_file_to_signed_url,
     zip_files,
@@ -23,6 +22,7 @@
 app = typer.Typer()
 manager_app = typer.Typer()
 workspace_manager = WorkspaceManager()
+registry_api = RegistryAPI()
 
 
 def execute_cm_cli(args, channel=None, mode=None):
@@ -534,7 +534,7 @@ def publish(
 
     # Call API to fetch node version with the token in the body
     typer.echo("Publishing node version...")
-    response = publish_node_version(config, token)
+    response = registry_api.publish_node_version(config, token)
 
     # Zip up all files in the current directory, respecting .gitignore files.
     signed_url = response.signedUrl

comfy_cli/registry/__init__.py

@@ -1,11 +1,11 @@
-from .api import publish_node_version, upload_file_to_signed_url
+from .api import RegistryAPI, upload_file_to_signed_url
 
 from .config_parser import extract_node_configuration, initialize_project_config
 from .types import PyProjectConfig, PublishNodeVersionResponse, NodeVersion
 from .zip import zip_files
 
 __all__ = [
-    "publish_node_version",
+    "RegistryAPI",
     "extract_node_configuration",
     "PyProjectConfig",
     "PublishNodeVersionResponse",

comfy_cli/registry/api.py

@@ -1,62 +1,72 @@
+import os
+
 import requests
 import json
-from comfy_cli import constants
 from comfy_cli.registry.types import (
     PyProjectConfig,
     PublishNodeVersionResponse,
     NodeVersion,
 )
 
 
-def publish_node_version(
-    node_config: PyProjectConfig, token: str
-) -> PublishNodeVersionResponse:
-    """
-    Publishes a new version of a node.
+class RegistryAPI:
+    def __init__(self):
+        self.base_url = self.determine_base_url()
+
+    def determine_base_url(self):
+        if os.getenv("ENVIRONMENT") == "dev":
+            return "http://localhost:8080"
+        else:
+            return "https://api-frontend-dev-qod3oz2v2q-uc.a.run.app"
+
+    def publish_node_version(
+        self, node_config: PyProjectConfig, token: str
+    ) -> PublishNodeVersionResponse:
+        """
+        Publishes a new version of a node.
 
-    Args:
-    node_config (PyProjectConfig): The node configuration.
-    token (str): Personal access token for authentication.
+        Args:
+          node_config (PyProjectConfig): The node configuration.
+          token (str): The token to authenticate with the API server.
 
-    Returns:
-    dict: JSON response from the API server.
-    """
-    url = f"{constants.COMFY_REGISTRY_URL_ROOT}/publishers/{node_config.tool_comfy.publisher_id}/nodes/{node_config.project.name}/versions"
-    headers = {"Content-Type": "application/json"}
-    body = {
-        "personal_access_token": token,
-        "node": {
-            "id": node_config.project.name,
-            "description": node_config.project.description,
-            "name": node_config.tool_comfy.display_name,
-            "license": node_config.project.license,
-            "repository": node_config.project.urls.repository,
-        },
-        "node_version": {
-            "version": node_config.project.version,
-            "dependencies": node_config.project.dependencies,
-        },
-    }
+        Returns:
+        PublishNodeVersionResponse: The response object from the API server.
+        """
+        url = f"{self.base_url}/publishers/{node_config.tool_comfy.publisher_id}/nodes/{node_config.project.name}/versions"
+        headers = {"Content-Type": "application/json"}
+        body = {
+            "personal_access_token": token,
+            "node": {
+                "id": node_config.project.name,
+                "description": node_config.project.description,
+                "name": node_config.tool_comfy.display_name,
+                "license": node_config.project.license,
+                "repository": node_config.project.urls.repository,
+            },
+            "node_version": {
+                "version": node_config.project.version,
+                "dependencies": node_config.project.dependencies,
+            },
+        }
 
-    response = requests.post(url, headers=headers, data=json.dumps(body))
+        response = requests.post(url, headers=headers, data=json.dumps(body))
 
-    if response.status_code == 201:
-        data = response.json()
-        node_version = NodeVersion(
-            changelog=data["node_version"]["changelog"],
-            dependencies=data["node_version"]["dependencies"],
-            deprecated=data["node_version"]["deprecated"],
-            id=data["node_version"]["id"],
-            version=data["node_version"]["version"],
-        )
-        node_data = PublishNodeVersionResponse(
-            node_version=node_version, signedUrl=data["signedUrl"]
-        )
-        return node_data
-    else:
-        raise Exception(
-            f"Failed to publish node version: {response.status_code} {response.text}"
-        )
+        if response.status_code == 201:
+            data = response.json()
+            node_version = NodeVersion(
+                changelog=data["node_version"]["changelog"],
+                dependencies=data["node_version"]["dependencies"],
+                deprecated=data["node_version"]["deprecated"],
+                id=data["node_version"]["id"],
+                version=data["node_version"]["version"],
+            )
+            return PublishNodeVersionResponse(
+                node_version=node_version, signedUrl=data["signedUrl"]
+            )
+        else:
+            raise Exception(
+                f"Failed to publish node version: {response.status_code} {response.text}"
+            )
 
 
 def upload_file_to_signed_url(signed_url: str, file_path: str):