diff --git a/comfy_cli/command/custom_nodes/command.py b/comfy_cli/command/custom_nodes/command.py index f15d6072..d826c6c3 100644 --- a/comfy_cli/command/custom_nodes/command.py +++ b/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.") diff --git a/comfy_cli/file_utils.py b/comfy_cli/file_utils.py index 7cc2abf6..640d462c 100644 --- a/comfy_cli/file_utils.py +++ b/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): diff --git a/comfy_cli/registry/config_parser.py b/comfy_cli/registry/config_parser.py index db904407..a1120381 100644 --- a/comfy_cli/registry/config_parser.py +++ b/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) diff --git a/comfy_cli/registry/types.py b/comfy_cli/registry/types.py index 0c2b7eb1..b3f94b2b 100644 --- a/comfy_cli/registry/types.py +++ b/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 diff --git a/tests/comfy_cli/command/nodes/test_publish.py b/tests/comfy_cli/command/nodes/test_publish.py index ccc58c5f..250d2966 100644 --- a/tests/comfy_cli/command/nodes/test_publish.py +++ b/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