Added: publish all files in includes field in pyproject.toml (#270)

Author: jtydhr88

comfy_cli/command/custom_nodes/command.py

@@ -765,7 +765,13 @@ def publish(
         signed_url = response.signedUrl
         zip_filename = NODE_ZIP_FILENAME
         typer.echo("Creating zip file...")
-        zip_files(zip_filename)
+
+        includes = config.tool_comfy.includes if config and config.tool_comfy else []
+
+        if includes:
+            typer.echo(f"Including additional directories: {', '.join(includes)}")
+
+        zip_files(zip_filename, includes=includes)
 
         # Upload the zip file to the signed URL
         typer.echo("Uploading zip file...")
@@ -917,7 +923,13 @@ def pack():
         raise typer.Exit(code=1)
 
     zip_filename = NODE_ZIP_FILENAME
-    zip_files(zip_filename)
+    includes = config.tool_comfy.includes if config and config.tool_comfy else []
+
+    if includes:
+        typer.echo(f"Including additional directories: {', '.join(includes)}")
+
+    zip_files(zip_filename, includes=includes)
+
     typer.echo(f"Created zip file: {NODE_ZIP_FILENAME}")
     logging.info("Node has been packed successfully.")
 

comfy_cli/file_utils.py

@@ -88,39 +88,61 @@ def download_file(url: str, local_filepath: pathlib.Path, headers: Optional[dict
             raise DownloadException(f"Failed to download file.\n{status_reason}")
 
 
-def zip_files(zip_filename):
+def zip_files(zip_filename, includes=None):
     """
-    Zip all files in the current directory that are tracked by git.
+    Zip all files in the current directory that are tracked by git,
+    plus any additional directories specified in includes.
     """
+    includes = includes or []
+    included_paths = set()
+    git_files = []
+
     try:
-        # Get list of git-tracked files using git ls-files
         import subprocess
 
         git_files = subprocess.check_output(["git", "ls-files"], text=True).splitlines()
-        # Zip only git-tracked files
-        with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
+    except (subprocess.SubprocessError, FileNotFoundError):
+        print("Warning: Not in a git repository or git not installed. Zipping all files.")
+
+    # Zip only git-tracked files
+    with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
+        if git_files:
             for file_path in git_files:
                 if zip_filename in file_path:
                     continue
                 if os.path.exists(file_path):
                     zipf.write(file_path)
+                    included_paths.add(file_path)
                 else:
                     print(f"File not found. Not including in zip: {file_path}")
-        return
-    except (subprocess.SubprocessError, FileNotFoundError):
-        print("Warning: Not in a git repository or git not installed. Zipping all files.")
-
-    with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
-        for root, dirs, files in os.walk("."):
-            if ".git" in dirs:
-                dirs.remove(".git")
-            for file in files:
-                file_path = os.path.join(root, file)
-                # Skip zipping the zip file itself
-                if zip_filename in file_path:
-                    continue
-                relative_path = os.path.relpath(file_path, start=".")
-                zipf.write(file_path, relative_path)
+        else:
+            for root, dirs, files in os.walk("."):
+                if ".git" in dirs:
+                    dirs.remove(".git")
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    # Skip zipping the zip file itself
+                    if zip_filename in file_path:
+                        continue
+                    relative_path = os.path.relpath(file_path, start=".")
+                    zipf.write(file_path, relative_path)
+                    included_paths.add(file_path)
+
+        for include_dir in includes:
+            include_dir = include_dir.lstrip("/")
+            if not os.path.exists(include_dir):
+                print(f"Warning: Included directory '{include_dir}' does not exist, creating empty directory")
+                zipf.writestr(f"{include_dir}/", "")
+                continue
+
+            for root, dirs, files in os.walk(include_dir):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    if zip_filename in file_path or file_path in included_paths:
+                        continue
+                    relative_path = os.path.relpath(file_path, start=".")
+                    zipf.write(file_path, relative_path)
+                    included_paths.add(file_path)
 
 
 def upload_file_to_signed_url(signed_url: str, file_path: str):

comfy_cli/registry/config_parser.py

@@ -42,6 +42,7 @@ def create_comfynode_config():
     comfy["PublisherId"] = ""
     comfy["DisplayName"] = "ComfyUI-AIT"
     comfy["Icon"] = ""
+    comfy["includes"] = tomlkit.array()
 
     tool.add("comfy", comfy)
     document.add("tool", tool)
@@ -196,6 +197,7 @@ def extract_node_configuration(
         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", []),
     )
 
     return PyProjectConfig(project=project, tool_comfy=comfy)

comfy_cli/registry/types.py

@@ -51,6 +51,7 @@ class ComfyConfig:
     display_name: str = ""
     icon: str = ""
     models: List[Model] = field(default_factory=list)
+    includes: List[str] = field(default_factory=list)
 
 
 @dataclass

tests/comfy_cli/command/nodes/test_publish.py

@@ -7,6 +7,29 @@
 runner = CliRunner()
 
 
+def create_mock_config(includes_list=None):
+    if includes_list is None:
+        includes_list = []
+
+    mock_pyproject_config = MagicMock()
+
+    mock_tool_comfy_section = MagicMock()
+    mock_tool_comfy_section.name = "test-node"
+    mock_tool_comfy_section.version = "0.1.0"
+    mock_tool_comfy_section.description = "A test node."
+    mock_tool_comfy_section.author = "Test Author"
+    mock_tool_comfy_section.license = "MIT"
+    mock_tool_comfy_section.tags = ["test"]
+    mock_tool_comfy_section.repository = "http://example.com/repo"
+    mock_tool_comfy_section.homepage = "http://example.com/home"
+    mock_tool_comfy_section.documentation = "http://example.com/docs"
+    mock_tool_comfy_section.includes = includes_list
+
+    mock_pyproject_config.tool_comfy = mock_tool_comfy_section
+
+    return mock_pyproject_config
+
+
 def test_publish_fails_on_security_violations():
     # Mock subprocess.run to simulate security violations
     mock_result = MagicMock()
@@ -40,7 +63,8 @@ def test_publish_continues_on_no_security_violations():
         patch("comfy_cli.command.custom_nodes.command.upload_file_to_signed_url") as mock_upload,
     ):
         # Setup the mocks
-        mock_extract.return_value = {"name": "test-node"}
+        mock_extract.return_value = create_mock_config()
+
         mock_prompt.return_value = "test-token"
         mock_publish.return_value = MagicMock(signedUrl="https://test.url")
 
@@ -76,7 +100,8 @@ def test_publish_with_token_option():
         patch("comfy_cli.command.custom_nodes.command.upload_file_to_signed_url") as mock_upload,
     ):
         # Setup the mocks
-        mock_extract.return_value = {"name": "test-node"}
+        mock_extract.return_value = create_mock_config()
+
         mock_publish.return_value = MagicMock(signedUrl="https://test.url")
 
         # Run the publish command with token
@@ -104,7 +129,8 @@ def test_publish_exits_on_upload_failure():
         patch("comfy_cli.command.custom_nodes.command.upload_file_to_signed_url") as mock_upload,
     ):
         # Setup the mocks
-        mock_extract.return_value = {"name": "test-node"}
+        mock_extract.return_value = create_mock_config()
+
         mock_publish.return_value = MagicMock(signedUrl="https://test.url")
         mock_upload.side_effect = Exception("Upload failed with status code: 403")
 
@@ -117,3 +143,33 @@ def test_publish_exits_on_upload_failure():
         assert mock_publish.called
         assert mock_zip.called
         assert mock_upload.called
+
+
+def test_publish_with_includes_parameter():
+    # Mock subprocess.run to simulate no violations
+    mock_result = MagicMock()
+    mock_result.returncode = 0
+    mock_result.stdout = ""
+
+    with (
+        patch("subprocess.run", return_value=mock_result),
+        patch("comfy_cli.command.custom_nodes.command.extract_node_configuration") as mock_extract,
+        patch("comfy_cli.command.custom_nodes.command.registry_api.publish_node_version") as mock_publish,
+        patch("comfy_cli.command.custom_nodes.command.zip_files") as mock_zip,
+        patch("comfy_cli.command.custom_nodes.command.upload_file_to_signed_url") as mock_upload,
+    ):
+        includes = ["/js", "/dist"]
+
+        # Setup the mocks
+        mock_extract.return_value = create_mock_config(includes)
+
+        mock_publish.return_value = MagicMock(signedUrl="https://test.url")
+
+        # Run the publish command with token
+        _result = runner.invoke(app, ["publish", "--token", "test-token"])
+
+        # Verify the publish flow worked with provided token
+        assert mock_extract.called
+        assert mock_publish.called
+        assert mock_zip.called
+        assert mock_upload.called