From f74f091c015d00b3dc85cfeef0a0f32f01f08e10 Mon Sep 17 00:00:00 2001 From: HowToSD Date: Thu, 13 Feb 2025 23:38:27 -0800 Subject: [PATCH] Added Pandas Save CSV node --- README.md | 3 +- __init__.py | 4 ++ docs/ug.md | 3 ++ modules/pandas/pandas_save_csv.py | 63 +++++++++++++++++++++++++++++++ pyproject.toml | 2 +- web/js/SaveCSV.js | 36 ++++++++++++++++++ 6 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 modules/pandas/pandas_save_csv.py create mode 100644 web/js/SaveCSV.js diff --git a/README.md b/README.md index c8df137..b5b3a01 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,8 @@ ComfyUI-Data-Analysis is a custom module for analyzing data using Pandas and Mat Currently, it supports the following nodes: -- Loading CSV as a DataFrame +- Loading a CSV file as a DataFrame +- Saving a DataFrame to a CSV file - Selecting a subset of columns from a DataFrame - Selecting a subset of rows from a DataFrame by specifying a filter condition - Sorting a DataFrame by a column diff --git a/__init__.py b/__init__.py index 33af2a5..3eb45d9 100644 --- a/__init__.py +++ b/__init__.py @@ -5,6 +5,7 @@ from .modules.matplotlib.mpl_scatter import MPLScatter from .modules.pandas.pandas_head import PandasHead from .modules.pandas.pandas_join import PandasJoin from .modules.pandas.pandas_load_csv import PandasLoadCSV +from .modules.pandas.pandas_save_csv import PandasSaveCSV from .modules.pandas.pandas_select_columns import PandasSelectColumns from .modules.pandas.pandas_select_rows import PandasSelectRows from .modules.pandas.pandas_show_dataframe import PandasShowDataFrame @@ -21,6 +22,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PandasHead": PandasHead, "PandasJoin": PandasJoin, "PandasLoadCSV": PandasLoadCSV, + "PandasSaveCSV": PandasSaveCSV, "PandasSelectColumns": PandasSelectColumns, "PandasSelectRows": PandasSelectRows, "PandasShowDataFrame": PandasShowDataFrame, @@ -38,6 +40,7 @@ NODE_CLASS_MAPPINGS (Dict[str, Type[T]]): - "PandasHead": Handles retrieving the first few rows of a DataFrame. - "PandasJoin": Provides functionality to join multiple DataFrames. - "PandasLoadCSV": Loads CSV data into a DataFrame. + - "PandasSaveCSV": Save a DataFrame to a CSV file. - "PandasSelectColumns": Selects specific columns from a DataFrame. - "PandasSelectRows": Filters specific rows from a DataFrame. - "PandasShowDataFrame": Displays the content of the DataFrame. @@ -53,6 +56,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PandasHead": "Pandas Head", "PandasJoin": "Pandas Join", "PandasLoadCSV": "Pandas Load CSV", + "PandasSaveCSV": "Pandas Save CSV", "PandasSelectColumns": "Pandas Select Columns", "PandasSelectRows": "Pandas Select Rows", "PandasShowDataFrame": "Pandas Show DataFrame", diff --git a/docs/ug.md b/docs/ug.md index 77f4a59..ca7f496 100644 --- a/docs/ug.md +++ b/docs/ug.md @@ -35,6 +35,9 @@ Use **Pandas Load CSV** node to load a CSV. This will emit the DataFrame in a JSON string. CSV file path is relative to the ComfyUI installation directory unless you specify the absolute file path. +## Saving data +Use **Pandas Save CSV** node to save the DataFrame to a CSV file. + ## Selecting columns and rows Use **Pandas Select Columns** for selecting columns, **Pandas Select Rows** to select rows by a filter condition. diff --git a/modules/pandas/pandas_save_csv.py b/modules/pandas/pandas_save_csv.py new file mode 100644 index 0000000..ca6176c --- /dev/null +++ b/modules/pandas/pandas_save_csv.py @@ -0,0 +1,63 @@ +""" +This code is based on +https://github.com/pythongosssss/ComfyUI-Custom-Scripts/blob/main/py/show_text.py + +See credit/credit.md for the full license. +""" +from typing import Any, Dict +from io import StringIO +import pandas as pd + +class PandasSaveCSV: + """ + PandasSaveCSV: + A class for saving a PandasDataFrame to a CSV file. + """ + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "dataframe_json": ("STRING", {"forceInput": True}), + "file_path": ("STRING", {"multiline": False}) + }, + "hidden": { + "unique_id": "UNIQUE_ID", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + } + + INPUT_IS_LIST = True + RETURN_TYPES = ("STRING",) + FUNCTION = "save_csv" + OUTPUT_NODE = True + OUTPUT_IS_LIST = (True,) + + CATEGORY = "Data Analysis" + + def save_csv(self, dataframe_json,file_path=None, unique_id=None, extra_pnginfo=None): + text = None + if unique_id is not None and extra_pnginfo is not None: + if not isinstance(extra_pnginfo, list): + print("Error: extra_pnginfo is not a list") + elif ( + not isinstance(extra_pnginfo[0], dict) + or "workflow" not in extra_pnginfo[0] + ): + print("Error: extra_pnginfo[0] is not a dict or missing 'workflow' key") + elif file_path is None or len(file_path) == 0: + print("Error: file path is not specified") + else: + # Deserialize JSON string to DataFrame + df = pd.read_json(StringIO(dataframe_json[0])) + df.to_csv(path_or_buf=file_path[0]) # Save as CSV + text = [str(df)] # wrap in list + workflow = extra_pnginfo[0]["workflow"] + node = next( + (x for x in workflow["nodes"] if str(x["id"]) == str(unique_id[0])), + None, + ) + if node: + node["widgets_values"] = [text] + else: + print("unique_id or extra_pnginfo is None.") + return {"ui": {"text": text}, "result": (text,)} diff --git a/pyproject.toml b/pyproject.toml index f7f0965..85a4f96 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "data-analysis" description = "Data analysis custom modules for ComfyUI - Use Pandas & Matplotlib within ComfyUI" -version = "1.0.4" +version = "1.0.5" license = {file = "LICENSE"} [project.urls] diff --git a/web/js/SaveCSV.js b/web/js/SaveCSV.js new file mode 100644 index 0000000..e6f54fc --- /dev/null +++ b/web/js/SaveCSV.js @@ -0,0 +1,36 @@ +/* +This code is based on +https://github.com/pythongosssss/ComfyUI-Custom-Scripts/blob/main/web/js/showText.js + +See credit/credit.md for the full license. +*/ + +import { app } from "../../../scripts/app.js"; +import { ComfyWidgets } from "../../../scripts/widgets.js"; + +// Displays input text on a node +app.registerExtension({ + name: "HowToSD.SaveCSV", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === "PandasSaveCSV") { // This needs to match NODE_CLASS_MAPPINGS in Python + function populate(text) { + // No op as we are not updating the node's UI for now. + } + + // When the node is executed we will be sent the input text, display this in the widget + const onExecuted = nodeType.prototype.onExecuted; + nodeType.prototype.onExecuted = function (message) { + onExecuted?.apply(this, arguments); + populate.call(this, message.text); + }; + + const onConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function () { + onConfigure?.apply(this, arguments); + if (this.widgets_values?.length) { + populate.call(this, this.widgets_values.slice(+this.widgets_values.length > 1)); + } + }; + } + }, +}); \ No newline at end of file