From 27b042abf82314c2e2cf13960701be5d03baef28 Mon Sep 17 00:00:00 2001 From: HowToSD Date: Sat, 22 Feb 2025 19:41:31 -0800 Subject: [PATCH] Added support for integer column label in Pandas Select Columns node. --- .../pandas_wrapper/pandas_select_columns.py | 10 +-- modules/pandas_wrapper/utils.py | 42 +++++++++++- pyproject.toml | 2 +- .../utils_column_labels_string_to_list.py | 66 +++++++++++++++++++ 4 files changed, 113 insertions(+), 7 deletions(-) create mode 100644 tests/pandas_wrapper/utils_column_labels_string_to_list.py diff --git a/modules/pandas_wrapper/pandas_select_columns.py b/modules/pandas_wrapper/pandas_select_columns.py index eee1775..baa98eb 100644 --- a/modules/pandas_wrapper/pandas_select_columns.py +++ b/modules/pandas_wrapper/pandas_select_columns.py @@ -1,6 +1,6 @@ from typing import Any, Dict -from io import StringIO import pandas as pd +from .utils import column_labels_string_to_list class PandasSelectColumns: """ @@ -8,7 +8,7 @@ class PandasSelectColumns: category: Data subset selection """ - + @classmethod def INPUT_TYPES(cls) -> Dict[str, Any]: """ @@ -25,10 +25,10 @@ class PandasSelectColumns: } RETURN_TYPES: tuple = ("DATAFRAME",) - FUNCTION: str = "select_columns" + FUNCTION: str = "f" CATEGORY: str = "Data Analysis" - def select_columns(self, dataframe: pd.DataFrame, column_names: str) -> tuple: + def f(self, dataframe: pd.DataFrame, column_names: str) -> tuple: """ Selects specific columns from a pandas DataFrame. @@ -39,6 +39,6 @@ class PandasSelectColumns: Returns: tuple: A tuple containing the DataFrame with selected columns. """ - selected_columns = [col.strip() for col in column_names.split(",")] + selected_columns = column_labels_string_to_list(dataframe, column_names) df2 = dataframe[selected_columns] return (df2,) diff --git a/modules/pandas_wrapper/utils.py b/modules/pandas_wrapper/utils.py index 7d39ec7..d2a0c09 100644 --- a/modules/pandas_wrapper/utils.py +++ b/modules/pandas_wrapper/utils.py @@ -1,11 +1,51 @@ +from typing import List, Union import json import numpy as np import pandas as pd +import pandas as pd + + +def column_labels_string_to_list(dataframe: pd.DataFrame, s: str) -> List[Union[str, int]]: + """ + Converts a comma-separated string of column labels into a list of valid column names for a Pandas DataFrame. + + This function ensures that each label in the input string corresponds to an existing column in the DataFrame. + If a label is a numeric string and corresponds to an integer column name, it will be converted to an integer. + If a column name does not exist in the DataFrame, a ValueError is raised. + + Args: + dataframe (pd.DataFrame): The DataFrame to validate column names against. + s (str): A comma-separated string representing column labels. + + Returns: + List[Union[str, int]]: A list of valid column labels. + + Raises: + ValueError: If any label in `s` does not match an existing column name. + """ + columns = [col.strip() for col in s.split(",")] + valid_columns = set(dataframe.columns) # Use a set for faster lookups + output_columns = [] + + for c in columns: + if c in valid_columns: + output_columns.append(c) + else: + try: + numeric_c = int(c) + if numeric_c in valid_columns: + output_columns.append(numeric_c) + else: + raise ValueError(f"Column '{c}' not found in the DataFrame.") + except ValueError: + raise ValueError(f"Column '{c}' not found in the DataFrame.") + + return output_columns def series_to_jsons(series: pd.Series) -> str: """ - Convert a pandas Series to a JSON string. + Converts a pandas Series to a JSON string. Args: series (pd.Series): The pandas Series to convert. diff --git a/pyproject.toml b/pyproject.toml index 76678a8..f387b11 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.1.21" +version = "1.1.22" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pandas_wrapper/utils_column_labels_string_to_list.py b/tests/pandas_wrapper/utils_column_labels_string_to_list.py new file mode 100644 index 0000000..e6b0099 --- /dev/null +++ b/tests/pandas_wrapper/utils_column_labels_string_to_list.py @@ -0,0 +1,66 @@ +import os +import sys +import unittest +import pandas as pd +from typing import List, Union + +PROJECT_ROOT = os.path.realpath(os.path.join(__file__, "..", "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from pandas_wrapper.utils import column_labels_string_to_list + +class TestColumnLabelsStringToList(unittest.TestCase): + def setUp(self): + """Set up test data.""" + self.df = pd.DataFrame({ + "A": [1, 2, 3], + "B": [4, 5, 6], + "C": [7, 8, 9], + 10: [10, 11, 12], # Integer column name + 20: [13, 14, 15] # Another integer column name + }) + + def test_valid_column_strings(self): + """Test valid column name conversion from string to list.""" + self.assertEqual(column_labels_string_to_list(self.df, "A,B"), ["A", "B"]) + self.assertEqual(column_labels_string_to_list(self.df, "C"), ["C"]) + + def test_integer_column_labels(self): + """Test conversion when integer column names are referenced as strings.""" + self.assertEqual(column_labels_string_to_list(self.df, "10"), [10]) + self.assertEqual(column_labels_string_to_list(self.df, "20"), [20]) + self.assertEqual(column_labels_string_to_list(self.df, "10,20"), [10, 20]) + + def test_mixed_string_and_integer_columns(self): + """Test function with mixed string and integer column labels.""" + self.assertEqual(column_labels_string_to_list(self.df, "A,10"), ["A", 10]) + self.assertEqual(column_labels_string_to_list(self.df, "B,20"), ["B", 20]) + + def test_whitespace_handling(self): + """Test that spaces around column names are stripped.""" + self.assertEqual(column_labels_string_to_list(self.df, " A , B "), ["A", "B"]) + self.assertEqual(column_labels_string_to_list(self.df, " 10 , 20 "), [10, 20]) + + def test_non_existent_column(self): + """Test behavior when a non-existent column is provided.""" + with self.assertRaises(ValueError) as context: + column_labels_string_to_list(self.df, "D") + self.assertIn("Column 'D' not found in the DataFrame.", str(context.exception)) + + def test_mixed_valid_and_invalid_columns(self): + """Test that function raises ValueError when a mix of valid and invalid columns is provided.""" + with self.assertRaises(ValueError) as context: + column_labels_string_to_list(self.df, "A,D") + self.assertIn("Column 'D' not found in the DataFrame.", str(context.exception)) + + def test_invalid_integer_conversion(self): + """Test that a string that can't be converted to an integer raises ValueError.""" + with self.assertRaises(ValueError) as context: + column_labels_string_to_list(self.df, "A,xyz") + self.assertIn("Column 'xyz' not found in the DataFrame.", str(context.exception)) + + +if __name__ == "__main__": + unittest.main() +