Added support for integer column label in Pandas Select Columns node.

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,)

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.

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]

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()
+