_

Author: jeromedockes

skrub/_data_ops/_evaluation.py

@@ -427,18 +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 to
-    #   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. We do perform
-    #   such a check when a key is missing from the env to provide a better
-    #   message, but it is only used for the content of the message rather than
-    #   enforcing no extra keys ahead of time.
+    # - 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
@@ -531,9 +528,11 @@ def _uninitialized_variable_msg(error, data_op, environment):
     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 not k.startswith("_skrub_")}.difference(
-            var_names + choice_names
-        )
+        {
+            k
+            for k in environment.keys()
+            if isinstance(k, str) and not k.startswith("_skrub_")
+        }.difference(var_names + choice_names)
     )
 
     msg = (

skrub/_data_ops/tests/test_errors.py

@@ -676,13 +676,16 @@ def test_missing_var_message():
 
 @pytest.mark.skipif(sys.version_info < (3, 11), reason="no add_note")
 def test_missing_var_message_train_test_split():
-    X = (skrub.var("x", np.arange(20)) + skrub.var("a")).skb.mark_as_X() + skrub.var(
-        "b"
-    )
+    b = skrub.var("b")
+    X = (skrub.var("x", np.arange(20)) + skrub.var("a")).skb.mark_as_X() + b
     # 'z' is extra but not 'b', even though 'b' is not needed for collecting X
     # it still exists in the DataOp as a whole.
     with pytest.raises(KeyError, match=r"\['z'\]"):
         X.skb.train_test_split({"z": 0, "b": 1})
+    # check that we still get the correct error when the env contains int (ID)
+    # keys
+    with pytest.raises(KeyError, match=r"\['z'\]") as exc:
+        X.skb.train_test_split({"z": 0, b.skb.id: 1})
     with pytest.raises(KeyError) as exc:
         X.skb.train_test_split()
     full_msg = "\n".join(traceback.format_exception(exc.value, exc.value, exc.tb))