fix: apply the dfn Narwhals shortcut in pre= to any function, including with selectors and segments

pointblank/validate.py

@@ -16633,37 +16633,8 @@ def interrogate(
                         else list(data_tbl_step.columns)
                     )
 
-                    # Read the text of the preprocessing function
-                    pre_text = _pre_processing_funcs_to_str(validation.pre)
-
-                    # Determine if the preprocessing function is a lambda function; return a boolean
-                    is_lambda = re.match(r"^lambda", pre_text) is not None
-
-                    # If the preprocessing function is a lambda function, then check if there is
-                    # a keyword argument called `dfn` in the lamda signature; if so, that's a cue
-                    # to use a Narwhalified version of the table
-                    if is_lambda:
-                        # Get the signature of the lambda function
-                        sig = inspect.signature(validation.pre)
-
-                        # Check if the lambda function has a keyword argument called `dfn`
-                        if "dfn" in sig.parameters:
-                            # Convert the table to a Narwhals DataFrame
-                            data_tbl_step = nw.from_native(data_tbl_step)
-
-                            # Apply the preprocessing function to the table
-                            data_tbl_step = validation.pre(dfn=data_tbl_step)
-
-                            # Convert the table back to its original format
-                            data_tbl_step = nw.to_native(data_tbl_step)
-
-                        else:
-                            # Apply the preprocessing function to the table
-                            data_tbl_step = validation.pre(data_tbl_step)
-
-                    # If the preprocessing function is a function, apply it to the table
-                    elif isinstance(validation.pre, Callable):
-                        data_tbl_step = validation.pre(data_tbl_step)
+                    # Apply the preprocessing function to the table
+                    data_tbl_step = _apply_pre(validation.pre, data_tbl_step)
 
                     # After successful preprocessing, check dimensions and create notes
                     # Use get_row_count() and get_column_count() for compatibility
@@ -22495,7 +22466,7 @@ def _evaluate_column_exprs(self, validation_info):
                 if validation.pre is None:
                     table = self.data
                 else:
-                    table = validation.pre(self.data)
+                    table = _apply_pre(validation.pre, self.data)
 
                 # Get the columns from the table as a list
                 columns = list(table.columns)
@@ -22593,7 +22564,7 @@ def _evaluate_segments(self, validation_info):
                 if validation.pre is None:
                     table = self.data
                 else:
-                    table = validation.pre(self.data)
+                    table = _apply_pre(validation.pre, self.data)
 
                 # If the `segments` expression is a string, that string is taken as a column name
                 # for which segmentation should occur across unique values in the column
@@ -25266,6 +25237,26 @@ def _transform_assertion_str(
     return type_upd
 
 
+def _apply_pre(pre: Callable, tbl: Any) -> Any:
+    """
+    Apply a `pre=` function to a table.
+
+    A function with a parameter named `dfn` (a lambda, a named function, or any other callable) is
+    given the table as a Narwhals DataFrame, so that it can use Narwhals expressions on any backend;
+    its result is converted back to the native table type. Other functions get the native table.
+    """
+    try:
+        parameters = inspect.signature(pre).parameters
+    except (TypeError, ValueError):  # pragma: no cover
+        # Some built-in callables have no inspectable signature
+        parameters = {}
+
+    if "dfn" in parameters:
+        return nw.to_native(pre(dfn=nw.from_native(tbl)), pass_through=True)
+
+    return pre(tbl)
+
+
 def _pre_processing_funcs_to_str(pre: Callable) -> str | list[str] | None:
     if isinstance(pre, Callable):
         return _get_callable_source(fn=pre)

tests/test_validate.py

@@ -19383,6 +19383,59 @@ def pre_x_ge_3(df):
     assert n_values[2] == 2
 
 
+def _times_ten_dfn(dfn):
+    return dfn.with_columns(b=nw.col("a") * 10)
+
+
+class _TimesTenDfn:
+    def __call__(self, dfn):
+        return dfn.with_columns(b=nw.col("a") * 10)
+
+
+# Defined away from the method call (the source text no longer starts with `lambda`)
+_times_ten_dfn_lambda = lambda dfn: dfn.with_columns(b=nw.col("a") * 10)  # noqa: E731
+
+DFN_PRE_FUNCTIONS = {
+    "lambda": _times_ten_dfn_lambda,
+    "named": _times_ten_dfn,
+    "callable_object": _TimesTenDfn(),
+    "partial": partial(lambda factor, dfn: dfn.with_columns(b=nw.col("a") * factor), 10),
+}
+
+
+@pytest.mark.parametrize("tbl_type", ["polars", "pandas"])
+@pytest.mark.parametrize("kind", list(DFN_PRE_FUNCTIONS))
+def test_pre_dfn_any_callable(tbl_type, kind) -> None:
+    # Any `pre=` function with a `dfn` parameter gets a Narwhals table: with plain columns, with
+    # columns resolved by a selector, and with segments
+    tbl = pl.DataFrame({"a": [1, 2, 3, 4], "g": ["x", "x", "y", "y"]})
+    tbl = tbl if tbl_type == "polars" else tbl.to_pandas()
+    pre = DFN_PRE_FUNCTIONS[kind]
+
+    validation = (
+        Validate(tbl)
+        .col_vals_gt(columns="b", value=15, pre=pre)
+        .col_vals_gt(columns=col(starts_with("b")), value=15, pre=pre)
+        .col_vals_gt(columns="b", value=15, pre=pre, segments="g")
+        .interrogate()
+    )
+
+    assert not any(step.eval_error for step in validation.validation_info)
+    assert validation.n_passed(scalar=False) == {1: 3, 2: 3, 3: 1, 4: 2}
+    assert validation.n_failed(scalar=False) == {1: 1, 2: 1, 3: 1, 4: 0}
+
+
+def test_pre_dfn_result_is_native() -> None:
+    # The Narwhals table a `dfn` function returns is converted back to the native type
+    from pointblank.validate import _apply_pre
+
+    tbl = pl.DataFrame({"a": [1, 2]})
+
+    assert isinstance(_apply_pre(_times_ten_dfn, tbl), pl.DataFrame)
+    # A function that already returns a native table is fine too
+    assert isinstance(_apply_pre(lambda dfn: dfn.to_native(), tbl), pl.DataFrame)
+
+
 @pytest.mark.parametrize("tbl_fixture", ["tbl_pd", "tbl_pl"])
 def test_pre_parameter_native_lambda_isolation(request, tbl_fixture) -> None:
     """Test pre parameter isolation using native lambdas specific to each DataFrame type."""

user_guide/01-validation-plan/04-preprocessing.qmd

@@ -231,9 +231,16 @@ have to go back to a Pandas DataFrame by using `.to_native()`). Pointblank is ab
 with the Narwhals DataFrame for validation purposes, which makes the workflow more concise.
 
 One more thing to note: Pointblank provides a convenient syntactic sugar for working with Narwhals.
-If you name the lambda parameter `dfn` instead of `df`, the system automatically applies
-`nw.from_native()` to the input DataFrame first. This lets you write more concise code without
-having to explicitly convert the DataFrame to a Narwhals format.
+If you name the parameter of the preprocessing function (a lambda or a regular function) `dfn`
+instead of `df`, the system automatically applies `nw.from_native()` to the input DataFrame first,
+and converts the result back to the original DataFrame type. This lets you write more concise code
+without having to explicitly convert the DataFrame to a Narwhals format. The function above could
+be written like this:
+
+```python
+def get_median_columns_c_and_a(dfn):
+    return dfn.select(nw.median("c"), nw.median("a"))
+```
 
 ## Swapping in a Totally Different DataFrame