expose tags, classes_, _estimator_type in Learner classes

Author: jeromedockes

skrub/_data_ops/_estimator.py

@@ -82,16 +82,57 @@ def _copy_attr(source, target, attributes):
             pass
 
 
-class _CloudPickleDataOp(_CloudPickle):
+class _DataOpWrapperMixin(_CloudPickle):
     """
-    Mixin to serialize the `DataOp` attribute with cloudpickle when pickling a
-    learner.
+    Mixin for learners and estimators that wrap a DataOp.
+
+    It exposes some attributes (tags, classes_) needed by scikit-learn by
+    inspecting the data_op attribute.
+
+    It also relies on the _CloudPickle mixin to serialize the DataOp with
+    cloudpickle rather than the normal pickle protocol when the learner is
+    pickled.
     """
 
     _cloudpickle_attributes = ["data_op"]
 
+    @property
+    def _estimator_type(self):
+        first = find_first_apply(self.data_op)
+        if first is None:
+            return "transformer"
+        estimator = get_default(first._skrub_impl.estimator)
+        if isinstance(estimator, DataOp):
+            return "transformer"
+        try:
+            return estimator._estimator_type
+        except AttributeError:
+            return "transformer"
+
+    if hasattr(BaseEstimator, "__sklearn_tags__"):
+        # scikit-learn >= 1.6
+
+        def __sklearn_tags__(self):
+            first = find_first_apply(self.data_op)
+            if first is None:
+                return _default_sklearn_tags()
+            estimator = get_default(first._skrub_impl.estimator)
+            if isinstance(estimator, DataOp):
+                return _default_sklearn_tags()
+            try:
+                return estimator.__sklearn_tags__()
+            except AttributeError:
+                return _default_sklearn_tags()
+
+    @property
+    def classes_(self):
+        try:
+            return _get_classes(self.data_op)
+        except AttributeError:
+            attribute_error(self, "classes_")
+
 
-class SkrubLearner(_CloudPickleDataOp, BaseEstimator):
+class SkrubLearner(_DataOpWrapperMixin, BaseEstimator):
     """Learner that evaluates a skrub DataOp.
 
     This class is not meant to be instantiated manually, ``SkrubLearner``
@@ -482,41 +523,6 @@ def _get_env(self, X, y):
             xy_environment[Y_NAME] = y
         return {**self.environment, **xy_environment}
 
-    @property
-    def _estimator_type(self):
-        first = find_first_apply(self.data_op)
-        if first is None:
-            return "transformer"
-        estimator = get_default(first._skrub_impl.estimator)
-        if isinstance(estimator, DataOp):
-            return "transformer"
-        try:
-            return estimator._estimator_type
-        except AttributeError:
-            return "transformer"
-
-    if hasattr(BaseEstimator, "__sklearn_tags__"):
-        # scikit-learn >= 1.6
-
-        def __sklearn_tags__(self):
-            first = find_first_apply(self.data_op)
-            if first is None:
-                return _default_sklearn_tags()
-            estimator = get_default(first._skrub_impl.estimator)
-            if isinstance(estimator, DataOp):
-                return _default_sklearn_tags()
-            try:
-                return estimator.__sklearn_tags__()
-            except AttributeError:
-                return _default_sklearn_tags()
-
-    @property
-    def classes_(self):
-        try:
-            return _get_classes(self.data_op)
-        except AttributeError:
-            attribute_error(self, "classes_")
-
 
 class _XyPipeline(_XyPipelineMixin, SkrubLearner):
     """
@@ -763,7 +769,7 @@ def iter_cv_splits(data_op, environment, *, keep_subsampling=False, cv=KFOLD_5):
         yield split_info
 
 
-class _BaseParamSearch(_CloudPickleDataOp, BaseEstimator):
+class _BaseParamSearch(_DataOpWrapperMixin, BaseEstimator):
     """Base class for hyperparameter search objects.
 
     It defines some default implementations for getting results, plotting, and