check if missing method exists in .skb namespace (#1541)

Author: jeromedockes

skrub/_data_ops/_data_ops.py

@@ -1326,21 +1326,22 @@ def compute(self, e, mode, environment):
         return self.target._skrub_impl.estimator_
 
 
+def _checked_get_attr(obj, obj_value, attr_name):
+    try:
+        return getattr(obj_value, attr_name)
+    except AttributeError:
+        if isinstance(obj, DataOp) and hasattr(obj.skb, attr_name):
+            comment = f"Did you mean `.skb.{attr_name}` and forget the `.skb`?"
+        else:
+            comment = None
+        attribute_error(obj_value, attr_name, comment)
+
+
 class GetAttr(DataOpImpl):
     _fields = ["source_object", "attr_name"]
 
     def compute(self, e, mode, environment):
-        try:
-            return getattr(e.source_object, e.attr_name)
-        except AttributeError:
-            pass
-        if isinstance(self.source_object, DataOp) and hasattr(
-            self.source_object.skb, e.attr_name
-        ):
-            comment = f"Did you mean `.skb.{e.attr_name}` and forget the `.skb`?"
-        else:
-            comment = None
-        attribute_error(e.source_object, e.attr_name, comment)
+        return _checked_get_attr(self.source_object, e.source_object, e.attr_name)
 
     def __repr__(self):
         return f"<{self.__class__.__name__} {short_repr(self.attr_name)}>"
@@ -1421,8 +1422,9 @@ class CallMethod(DataOpImpl):
     _fields = ["obj", "method_name", "args", "kwargs"]
 
     def compute(self, e, mode, environment):
+        method = _checked_get_attr(self.obj, e.obj, e.method_name)
         try:
-            return getattr(e.obj, e.method_name)(*e.args, **e.kwargs)
+            return method(*e.args, **e.kwargs)
         except Exception as err:
             # Better error message if we used the pandas DataFrame's `apply()`
             # but we meant `.skb.apply()`

skrub/_data_ops/tests/test_errors.py

@@ -1,5 +1,6 @@
 import pickle
 import re
+import traceback
 
 import numpy as np
 import pandas as pd
@@ -406,6 +407,28 @@ def test_attribute_errors():
         skrub.X().skb.applied_estimator
 
 
+def test_call_method_errors():
+    # In this case we have a preview value, so the access to `mark_as_X` fails
+    # right away, before we add the call (when creating a GetAttr node)
+    a = skrub.var("a", 0)
+    with pytest.raises(
+        (AttributeError, RuntimeError),
+        match=r"(?sm).*^Did you mean `\.skb\.mark_as_X`",
+    ):
+        a.mark_as_X()
+    a = skrub.var("a").mark_as_X()
+
+    # In this case we have no preview value so the error only happens when we
+    # call `.skb.eval()` and provide a value. By then the GetAttr + Call have
+    # been collapsed in the graph into a CallMethod node, so the error
+    # originates from a different place but it should provide the same info.
+    with pytest.raises((AttributeError, RuntimeError)) as exc:
+        a.skb.eval({"a": 0})
+    assert "Did you mean `.skb.mark_as_X` and forget the `.skb`?" in "\n".join(
+        traceback.format_exception(exc.value, exc.value, exc.tb)
+    )
+
+
 def test_concat_horizontal_numpy():
     a = skrub.var("a", skrub.datasets.toy_orders().orders)
     b = skrub.var("b", np.eye(3))