Validate CLI comma-separated column arguments (#361)

Summary: Pull Request resolved: #361

Reviewed By: sahil350

Differential Revision: D96316234

Pulled By: talgalili

fbshipit-source-id: 931ecabd7b6953f127aa21c3f7a96054f80ba2f4

Author: neuralsorcerer

CHANGELOG.md

@@ -79,6 +79,12 @@
     Series objects and returns `None` when both `below` and `above` are `None`,
     avoiding ambiguous concat inputs while preserving existing behavior for valid
     threshold sets.
+- **Validated and normalized comma-separated CLI column arguments**
+  - CLI column-list arguments now trim surrounding whitespace and reject empty
+    entries (for example, `"id,,weight"`) with clear `ValueError` messages,
+    preventing malformed column specifications from silently propagating.
+  - Applied to `--covariate_columns`, `--covariate_columns_for_diagnostics`,
+    `--batch_columns`, `--keep_columns`, and `--outcome_columns` parsing.
 
 ## Tests
 

balance/cli.py

@@ -12,7 +12,7 @@
 import logging
 from argparse import ArgumentParser, Namespace
 from pathlib import Path
-from typing import Any, Dict, List, Tuple, Type
+from typing import Any, Dict, List, Optional, Tuple, Type
 
 import balance
 import pandas as pd
@@ -25,6 +25,30 @@
 logger: logging.Logger = logging.getLogger(__package__)
 
 
+def _parse_csv_columns_arg(value: Optional[str], arg_name: str) -> List[str]:
+    """Parse a comma-separated CLI columns argument into a validated list.
+
+    Args:
+        value: Raw argparse string value.
+        arg_name: Argument name used for error context.
+
+    Returns:
+        A list of non-empty, trimmed column names.
+
+    Raises:
+        ValueError: If the value is missing or contains empty column names.
+    """
+    if value is None:
+        raise ValueError(f"{arg_name} cannot be None")
+
+    columns = [c.strip() for c in value.split(",")]
+    if not columns or any(c == "" for c in columns):
+        raise ValueError(
+            f"{arg_name} must be a comma-separated list of non-empty column names"
+        )
+    return columns
+
+
 class BalanceCLI:
     """Helper class that encapsulates CLI argument handling and execution.
 
@@ -195,7 +219,9 @@ def covariate_columns(self) -> List[str]:
                 BalanceCLI(Namespace(covariate_columns="x,y")).covariate_columns()
                 # ['x', 'y']
         """
-        return self.args.covariate_columns.split(",")
+        return _parse_csv_columns_arg(
+            self.args.covariate_columns, "--covariate_columns"
+        )
 
     def covariate_columns_for_diagnostics(self) -> List[str] | None:
         """Return covariate columns used for diagnostics reporting.
@@ -212,7 +238,11 @@ def covariate_columns_for_diagnostics(self) -> List[str] | None:
                 # ['x', 'y']
         """
         out = self.args.covariate_columns_for_diagnostics
-        return None if out is None else out.split(",")
+        return (
+            None
+            if out is None
+            else _parse_csv_columns_arg(out, "--covariate_columns_for_diagnostics")
+        )
 
     def rows_to_keep_for_diagnostics(self) -> str | None:
         """Return the diagnostics row-filter expression.
@@ -275,7 +305,7 @@ def batch_columns(self) -> List[str]:
                 BalanceCLI(Namespace(batch_columns="region,team")).batch_columns()
                 # ['region', 'team']
         """
-        return self.args.batch_columns.split(",")
+        return _parse_csv_columns_arg(self.args.batch_columns, "--batch_columns")
 
     def has_keep_columns(self) -> bool:
         """Return True when output keep columns are supplied.
@@ -318,8 +348,8 @@ def keep_columns(self) -> List[str] | None:
                 BalanceCLI(Namespace(keep_columns="id,weight")).keep_columns()
                 # ['id', 'weight']
         """
-        if self.args.keep_columns:
-            return self.args.keep_columns.split(",")
+        if self.args.keep_columns is not None:
+            return _parse_csv_columns_arg(self.args.keep_columns, "--keep_columns")
         return None
 
     def has_keep_row_column(self) -> bool:
@@ -376,8 +406,10 @@ def outcome_columns(self) -> List[str] | None:
                 BalanceCLI(Namespace(outcome_columns="y,z")).outcome_columns()
                 # ['y', 'z']
         """
-        if self.args.outcome_columns:
-            return self.args.outcome_columns.split(",")
+        if self.args.outcome_columns is not None:
+            return _parse_csv_columns_arg(
+                self.args.outcome_columns, "--outcome_columns"
+            )
         return None
 
     def max_de(self) -> float | None:

tests/test_cli.py

@@ -1446,6 +1446,32 @@ def test_keep_columns_returns_none_when_not_set(self) -> None:
         result = cli.keep_columns()
         self.assertIsNone(result)
 
+    def test_keep_columns_strips_whitespace(self) -> None:
+        """Test keep_columns trims whitespace around comma-separated names."""
+        args = Namespace(keep_columns=" id, weight ,extra ")
+        cli = BalanceCLI(args)
+        self.assertEqual(cli.keep_columns(), ["id", "weight", "extra"])
+
+    def test_keep_columns_raises_for_empty_column_name(self) -> None:
+        """Test keep_columns rejects empty names in comma-separated input."""
+        args = Namespace(keep_columns="id,,weight")
+        cli = BalanceCLI(args)
+        with self.assertRaisesRegex(
+            ValueError,
+            "--keep_columns must be a comma-separated list of non-empty column names",
+        ):
+            cli.keep_columns()
+
+    def test_keep_columns_raises_for_empty_string(self) -> None:
+        """Test keep_columns rejects an explicitly provided empty string."""
+        args = Namespace(keep_columns="")
+        cli = BalanceCLI(args)
+        with self.assertRaisesRegex(
+            ValueError,
+            "--keep_columns must be a comma-separated list of non-empty column names",
+        ):
+            cli.keep_columns()
+
     def test_has_keep_columns_with_keep_columns(self) -> None:
         """Test has_keep_columns returns True when keep_columns is set."""
         args = Namespace(keep_columns="id,weight")
@@ -1493,6 +1519,47 @@ def test_check_input_columns_raises_when_keep_column_missing(self) -> None:
         with self.assertRaises(AssertionError):
             cli.check_input_columns(columns)
 
+
+class TestBalanceCLI_csv_column_parsing(balance.testutil.BalanceTestCase):
+    """Test normalized parsing of comma-separated column arguments."""
+
+    def test_covariate_columns_strip_whitespace(self) -> None:
+        args = Namespace(covariate_columns=" a, b ,c ")
+        cli = BalanceCLI(args)
+        self.assertEqual(cli.covariate_columns(), ["a", "b", "c"])
+
+    def test_outcome_columns_raise_for_empty_name(self) -> None:
+        args = Namespace(outcome_columns="y,,z")
+        cli = BalanceCLI(args)
+        with self.assertRaisesRegex(
+            ValueError,
+            "--outcome_columns must be a comma-separated list of non-empty column names",
+        ):
+            cli.outcome_columns()
+
+    def test_outcome_columns_raise_for_empty_string(self) -> None:
+        args = Namespace(outcome_columns="")
+        cli = BalanceCLI(args)
+        with self.assertRaisesRegex(
+            ValueError,
+            "--outcome_columns must be a comma-separated list of non-empty column names",
+        ):
+            cli.outcome_columns()
+
+    def test_batch_columns_raise_for_empty_name(self) -> None:
+        args = Namespace(batch_columns="region,")
+        cli = BalanceCLI(args)
+        with self.assertRaisesRegex(
+            ValueError,
+            "--batch_columns must be a comma-separated list of non-empty column names",
+        ):
+            cli.batch_columns()
+
+    def test_covariate_columns_for_diagnostics_strip_whitespace(self) -> None:
+        args = Namespace(covariate_columns_for_diagnostics=" x, y ")
+        cli = BalanceCLI(args)
+        self.assertEqual(cli.covariate_columns_for_diagnostics(), ["x", "y"])
+
     def test_keep_columns_preserved_in_adjusted_output(self) -> None:
         """Test that --keep_columns columns survive adjustment via ignore_columns.