should report all missing columns rightaway

Author: vertti

daffy/decorators.py

@@ -13,19 +13,33 @@
 
 
 def _check_columns(df: DataFrameType, columns: ColumnsDef, strict: bool) -> None:
+    missing_columns = []
+    dtype_mismatches = []
+
     if isinstance(columns, list):
         for column in columns:
-            assert column in df.columns, f"Column {column} missing from DataFrame. Got {_describe_pd(df)}"
+            if column not in df.columns:
+                missing_columns.append(column)
     if isinstance(columns, dict):
         for column, dtype in columns.items():
-            assert column in df.columns, f"Column {column} missing from DataFrame. Got {_describe_pd(df)}"
-            assert df[column].dtype == dtype, (
-                f"Column {column} has wrong dtype. Was {df[column].dtype}, expected {dtype}"
-            )
-    if strict:
-        assert len(df.columns) == len(columns), (
-            f"DataFrame contained unexpected column(s): {', '.join(set(df.columns) - set(columns))}"
+            if column not in df.columns:
+                missing_columns.append(column)
+            elif df[column].dtype != dtype:
+                dtype_mismatches.append((column, df[column].dtype, dtype))
+
+    if missing_columns:
+        raise AssertionError(f"Missing columns: {missing_columns}. Got {_describe_pd(df)}")
+
+    if dtype_mismatches:
+        mismatches = ", ".join(
+            [f"Column {col} has wrong dtype. Was {was}, expected {expected}" for col, was, expected in dtype_mismatches]
         )
+        raise AssertionError(mismatches)
+
+    if strict:
+        extra_columns = set(df.columns) - set(columns)
+        if extra_columns:
+            raise AssertionError(f"DataFrame contained unexpected column(s): {', '.join(extra_columns)}")
 
 
 def df_out(columns: Optional[ColumnsDef] = None, strict: bool = False) -> Callable:

tests/test_decorators.py

@@ -107,7 +107,7 @@ def test_fn() -> DataFrameType:
     with pytest.raises(AssertionError) as excinfo:
         test_fn()
 
-    assert "Column FooColumn missing" in str(excinfo.value)
+    assert "Missing columns: ['FooColumn']. Got columns: ['Brand', 'Price']" in str(excinfo.value)
 
 
 def test_wrong_input_type_unnamed() -> None:
@@ -255,7 +255,17 @@ def test_fn(my_input: Any) -> Any:
 
     with pytest.raises(AssertionError) as excinfo:
         test_fn(df[["Brand"]])
-    assert "Column Price missing" in str(excinfo.value)
+    assert "Missing columns: ['Price']. Got columns: ['Brand']" in str(excinfo.value)
+
+@pytest.mark.parametrize(("df"), [pd.DataFrame(cars), pl.DataFrame(cars)])
+def test_df_in_missing_multiple_columns(df: DataFrameType) -> None:
+    @df_in(columns=["Brand", "Price", "Extra"])
+    def test_fn(my_input: Any) -> Any:
+        return my_input
+
+    with pytest.raises(AssertionError) as excinfo:
+        test_fn(df[["Brand"]])
+    assert "Missing columns: ['Price', 'Extra']. Got columns: ['Brand']" in str(excinfo.value)
 
 
 def test_df_out_with_df_modification(basic_pandas_df: pd.DataFrame, extended_pandas_df: pd.DataFrame) -> None: