better error message when variables are missing from env (#2211)

Author: jeromedockes

CHANGES.rst

@@ -16,6 +16,9 @@ New Features
 
 Changes
 -------
+- The error message when a key is missing from the environment passed to a
+  :class:`DataOp` or :class:`SkrubLearner` has been improved.
+  :pr:`2211` by :user:`Jérôme Dockès <jeromedockes>`.
 
 Bugfixes
 --------

skrub/_data_ops/_data_ops.py

@@ -150,6 +150,14 @@ class UninitializedVariable(KeyError):
     Evaluating a DataOp and a value has not been provided for one of the variables.
     """
 
+    def __init__(self, name):
+        self.name = name
+        self.message = f"No value has been provided for {name!r}"
+        super().__init__(name)
+
+    def __str__(self):
+        return self.message
+
 
 def _remove_shell_frames(stack):
     """
@@ -961,17 +969,15 @@ def compute(self, e, mode, environment):
         if mode == "preview":
             assert not environment
             if e.value is NULL:
-                raise UninitializedVariable(
-                    f"No value has been provided for {e.name!r}"
-                )
+                raise UninitializedVariable(e.name)
             return e.value
         if e.name in environment:
             return environment[e.name]
         if (
             e.becomes_default or environment.get("_skrub_use_var_values", False)
         ) and e.value is not NULL:
             return e.value
-        raise UninitializedVariable(f"No value has been provided for {e.name!r}")
+        raise UninitializedVariable(e.name)
 
     def preview_if_available(self):
         return self.value

skrub/_data_ops/_estimator.py

@@ -925,7 +925,10 @@ def _score(self, X, y=None, cast_to_float=True, return_predictions=False):
             return (result, {}) if return_predictions else result
         env = self._get_env(X, y)
         scorers = evaluate(
-            score_node._skrub_impl.scorers, mode="fit_transform", environment=env
+            score_node._skrub_impl.scorers,
+            mode="fit_transform",
+            environment=env,
+            ancestor_data_op=self.data_op,
         )
         all_scores = []
         cache = dict(env.get("_skrub_predictions", {}))
@@ -982,7 +985,13 @@ def _compute_X_y_and_cv(data_op, environment):
             # the estimator requests a y so some node must have been
             # marked as y
             raise ValueError('DataOp should have a node marked with "mark_as_y()"')
-    values = evaluate(nodes, mode="fit_transform", environment=environment, clear=True)
+    values = evaluate(
+        nodes,
+        mode="fit_transform",
+        environment=environment,
+        clear=True,
+        ancestor_data_op=data_op,
+    )
     if "y" in nodes:
         msg = (
             "\nAre `.skb.subsample()` and `.skb.mark_as_*()` applied in the same order"

skrub/_data_ops/_evaluation.py

@@ -26,10 +26,11 @@
     DataOp,
     Scoring,
     SplitX,
+    UninitializedVariable,
     Value,
     Var,
 )
-from ._utils import NULL, X_NAME, Y_NAME, simple_repr
+from ._utils import IS_PREVIEW_DATA_ENV_NAME, NULL, X_NAME, Y_NAME, simple_repr
 
 _BUILTIN_SEQ = (list, tuple, set, frozenset)
 
@@ -426,15 +427,15 @@ def _check_environment(environment):
         )
     # Notes about checking the env keys:
     #
-    # - env ⊂ variables: in some cases we could check that there are no extra
-    #   keys in `environment`, ie all keys in `environment` correspond to a
-    #   name in the DataOp. However in other cases we naturally end up
-    #   using a bigger environment than what is needed. For example we want tu
-    #   evaluate a sub-DataOp (such as the `mark_as_X()` node), and to do
-    #   it we use the environment that was passed to evaluate the full
-    #   DataOp. So if we want such a verification it should be a separate
-    #   check done at a higher level (eg in the estimators' `fit`, `predict`
-    #   etc.) where we know we are not working with a sub-DataOp.
+    # - env ⊂ variables: we could check that there are no extra keys in
+    #   `environment`, i.e. all keys in `environment` correspond to a name in
+    #   the DataOp. However in some cases we naturally end up using a bigger
+    #   environment than what is needed. For example we want to evaluate a
+    #   sub-DataOp (such as the result of `.skb.find()` or `.skb.find_X_y()`),
+    #   and to do it we use the environment created to evaluate the full
+    #   DataOp. We do perform this check when a key is missing from the env to
+    #   provide a better error message, but it is only used for the content of
+    #   the message rather than enforcing no extra keys ahead of time.
     #
     # - variables ⊂ env: we cannot check that all variables in the DataOp
     #   have a matching key in the `environment`, because depending on the
@@ -450,7 +451,14 @@ def _check_environment(environment):
 # consistent with .skb.eval() in evaluate's params
 
 
-def evaluate(data_op, mode="preview", environment=None, clear=False, callbacks=()):
+def evaluate(
+    data_op,
+    mode="preview",
+    environment=None,
+    clear=False,
+    callbacks=(),
+    ancestor_data_op=None,
+):
     """Evaluate a DataOp.
 
     Parameters
@@ -475,6 +483,14 @@ def evaluate(data_op, mode="preview", environment=None, clear=False, callbacks=(
         Each will be called, in the provided order, after evaluating each node.
         The signature is callback(data_op, result) where data_op is the DataOp
         that was just evaluated and result is the resulting value.
+
+    ancestor_data_op : DataOp or None
+        When we are evaluating a part of a DataOp (e.g. the X node only), pass
+        this so that we can look at all the variable names of the full DataOp
+        when producing error messages about missing or extra keys in the
+        environment. It is not used for other purposes than inspecting the
+        variable an choice names it contains. In particular it is not
+        evaluated.
     """
     _check_environment(environment)
     if clear:
@@ -486,11 +502,57 @@ def evaluate(data_op, mode="preview", environment=None, clear=False, callbacks=(
         return _Evaluator(mode=mode, environment=environment, callbacks=callbacks).run(
             data_op
         )
+    except UninitializedVariable as e:
+        if (
+            hasattr(e, "add_note")
+            and environment is not None
+            and not environment.get(IS_PREVIEW_DATA_ENV_NAME)
+        ):
+            # user passed an explicit environment rather than using the
+            # variables' preview values.
+            e.add_note(
+                _uninitialized_variable_msg(
+                    e,
+                    ancestor_data_op if ancestor_data_op is not None else data_op,
+                    environment,
+                )
+            )
+        raise
     finally:
         if clear:
             clear_results(data_op, mode=mode)
 
 
+def _uninitialized_variable_msg(error, data_op, environment):
+    missing_name = error.name
+    var_names = list(named_nodes(data_op).keys())
+    choice_names = [n for c in choices(data_op).values() if (n := c.name) is not None]
+    unused = list(
+        {
+            k
+            for k in environment.keys()
+            if isinstance(k, str) and not k.startswith("_skrub_")
+        }.difference(var_names + choice_names)
+    )
+
+    msg = (
+        "- Note that preview values passed to initialize skrub variables\n"
+        "  are ignored by default whenever we pass "
+        "an explicit 'environment' dictionary,\n"
+        "  for example when calling SkrubLearner.fit({'X': ..., 'y': ...}).\n"
+        f"  Please pass a value for {missing_name!r} in the environment.\n"
+        "  You can also use "
+        f"skrub.var({missing_name!r}, value=..., becomes_default=True)\n"
+        "  to always retain the initialization value as a default."
+    )
+    if unused:
+        msg += (
+            "\n- WARNING: the following keys were passed in the environment but "
+            f"have no corresponding variable in the DataOp:\n  {unused}"
+        )
+    return msg
+
+
 def _cache_pruner(data_op, mode):
     g = graph(data_op)
     ref_counts = {node_id: len(parents) for node_id, parents in g["parents"].items()}
@@ -675,6 +737,12 @@ def nodes(data_op):
     return list(graph(data_op)["nodes"].values())
 
 
+def named_nodes(data_op):
+    return {
+        name: op for op in nodes(data_op) if (name := op._skrub_impl.name) is not None
+    }
+
+
 def clear_results(data_op, mode=None):
     for n in nodes(data_op):
         if mode is None:

skrub/_data_ops/_skrub_namespace.py

@@ -42,6 +42,7 @@
     find_node_by_name,
     find_node_by_uuid,
     find_X_y_and_cv,
+    named_nodes,
     nodes,
 )
 from ._inspection import (
@@ -51,7 +52,7 @@
 )
 from ._optuna import OptunaParamSearch
 from ._subsampling import SubsamplePreviews, env_with_subsampling
-from ._utils import NULL, attribute_error
+from ._utils import IS_PREVIEW_DATA_ENV_NAME, NULL, attribute_error
 
 # By default, select all columns
 _SELECT_ALL_COLUMNS = s.all()
@@ -143,6 +144,11 @@ class SkrubNamespace:
     def __init__(self, data_op):
         self._data_op = data_op
 
+    def _get_env(self, environment=None):
+        if environment is not None:
+            return environment
+        return self.get_data() | {IS_PREVIEW_DATA_ENV_NAME: True}
+
     def _apply(
         self,
         estimator,
@@ -1121,7 +1127,7 @@ def eval(self, environment=None, *, keep_subsampling=False):
                 f"got: '{type(environment)}'"
             )
         if environment is None:
-            environment = self.get_data()
+            environment = self._get_env()
         else:
             environment = {
                 **environment,
@@ -1482,18 +1488,13 @@ def get_vars(self, all_named_ops=False):
         SkrubLearner(data_op=<BinOp: add>)
         """
         from ._data_ops import Var
-        from ._evaluation import nodes
 
-        named_nodes = {
-            name: op
-            for op in nodes(self._data_op)
-            if (name := op._skrub_impl.name) is not None
-        }
+        found_nodes = named_nodes(self._data_op)
         if all_named_ops:
-            return named_nodes
+            return found_nodes
         return {
             name: op
-            for name, op in named_nodes.items()
+            for name, op in found_nodes.items()
             if isinstance(op._skrub_impl, Var)
         }
 
@@ -1981,7 +1982,7 @@ def make_learner(self, *, fitted=False, keep_subsampling=False, choose="default"
         if not fitted:
             return learner
         return learner.fit(
-            env_with_subsampling(self._data_op, self.get_data(), keep_subsampling)
+            env_with_subsampling(self._data_op, self._get_env(), keep_subsampling)
         )
 
     @_check_before
@@ -2112,8 +2113,7 @@ def train_test_split(
                 category=FutureWarning,
             )
             split_func = splitter
-        if environment is None:
-            environment = self.get_data()
+        environment = self._get_env(environment)
         return train_test_split(
             self._data_op,
             environment,
@@ -2186,8 +2186,7 @@ def iter_cv_splits(self, environment=None, *, keep_subsampling=False, cv=None):
         >>> accuracies
         [1.0, 0.0, 1.0]
         """
-        if environment is None:
-            environment = self.get_data()
+        environment = self._get_env(environment)
         yield from iter_cv_splits(
             self._data_op, environment, keep_subsampling=keep_subsampling, cv=cv
         )
@@ -2301,7 +2300,7 @@ def make_grid_search(self, *, fitted=False, keep_subsampling=False, **kwargs):
         if not fitted:
             return search
         return search.fit(
-            env_with_subsampling(self._data_op, self.get_data(), keep_subsampling)
+            env_with_subsampling(self._data_op, self._get_env(), keep_subsampling)
         )
 
     @_check_before
@@ -2588,7 +2587,7 @@ def make_randomized_search(
         if not fitted:
             return search
         return search.fit(
-            env_with_subsampling(self._data_op, self.get_data(), keep_subsampling)
+            env_with_subsampling(self._data_op, self._get_env(), keep_subsampling)
         )
 
     @_check_before
@@ -2809,8 +2808,7 @@ def cross_validate(self, environment=None, *, keep_subsampling=False, **kwargs):
         4    0.90
         Name: test_score, dtype: float64
         """
-        if environment is None:
-            environment = self.get_data()
+        environment = self._get_env(environment)
 
         return cross_validate(
             self.make_learner(),

skrub/_data_ops/_utils.py

@@ -12,6 +12,7 @@
 FITTED_ESTIMATOR_METHODS = FITTED_PREDICTOR_METHODS + ("transform",)
 X_NAME = "_skrub_X"
 Y_NAME = "_skrub_y"
+IS_PREVIEW_DATA_ENV_NAME = "_skrub_is_preview_data_env"
 
 
 class Sentinels(enum.Enum):