FIX Fix available_if with bounded methods (scikit-learn#23077)

thomasjpfan · jjerphan · commit b51e8800f802 · 2022-04-29T14:06:40.000+02:00
diff --git a/doc/whats_new/v1.1.rst b/doc/whats_new/v1.1.rst
@@ -1071,6 +1071,9 @@ Changelog
 - |Fix| :func:`utils.estimator_html_repr` has an improved visualization for nested
   meta-estimators. :pr:`21310` by `Thomas Fan`_.
 
+- |Fix| :func:`utils.metaestimators.available_if` correctly returns a bounded
+  method that can be pickled. :pr:`23077` by `Thomas Fan`_.
+
 - |API| :func:`utils.metaestimators.if_delegate_has_method` is deprecated and will be
   removed in version 1.3. Use :func:`utils.metaestimators.available_if` instead.
   :pr:`22830` by :user:`Jérémie du Boisberranger <jeremiedbb>`.
diff --git a/sklearn/utils/metaestimators.py b/sklearn/utils/metaestimators.py
@@ -3,7 +3,9 @@
 #         Andreas Mueller
 # License: BSD
 from typing import List, Any
+from types import MethodType
 import warnings
+from functools import wraps
 
 from abc import ABCMeta, abstractmethod
 from operator import attrgetter
@@ -124,21 +126,17 @@ def __get__(self, obj, owner=None):
             # this is to allow access to the docstrings.
             if not self.check(obj):
                 raise attr_err
+            out = MethodType(self.fn, obj)
 
-            # lambda, but not partial, allows help() to work with update_wrapper
-            out = lambda *args, **kwargs: self.fn(obj, *args, **kwargs)  # noqa
         else:
-
-            def fn(*args, **kwargs):
+            # This makes it possible to use the decorated method as an unbound method,
+            # for instance when monkeypatching.
+            @wraps(self.fn)
+            def out(*args, **kwargs):
                 if not self.check(args[0]):
                     raise attr_err
                 return self.fn(*args, **kwargs)
 
-            # This makes it possible to use the decorated method as an unbound method,
-            # for instance when monkeypatching.
-            out = lambda *args, **kwargs: fn(*args, **kwargs)  # noqa
-        # update the docstring of the returned function
-        update_wrapper(out, self.fn)
         return out
 
 
diff --git a/sklearn/utils/tests/test_metaestimators.py b/sklearn/utils/tests/test_metaestimators.py
@@ -2,6 +2,8 @@
 import pytest
 import warnings
 
+import pickle
+
 from sklearn.utils.metaestimators import if_delegate_has_method
 from sklearn.utils.metaestimators import available_if
 
@@ -92,13 +94,14 @@ def test_if_delegate_has_method():
 class AvailableParameterEstimator:
     """This estimator's `available` parameter toggles the presence of a method"""
 
-    def __init__(self, available=True):
+    def __init__(self, available=True, return_value=1):
         self.available = available
+        self.return_value = return_value
 
     @available_if(lambda est: est.available)
     def available_func(self):
         """This is a mock available_if function"""
-        pass
+        return self.return_value
 
 
 def test_available_if_docstring():
@@ -155,3 +158,15 @@ def test_if_delegate_has_method_deprecated():
     # Only when calling it
     with pytest.warns(FutureWarning, match="if_delegate_has_method was deprecated"):
         hasattr(MetaEst(HasPredict()), "predict")
+
+
+def test_available_if_methods_can_be_pickled():
+    """Check that available_if methods can be pickled.
+
+    Non-regression test for #21344.
+    """
+    return_value = 10
+    est = AvailableParameterEstimator(available=True, return_value=return_value)
+    pickled_bytes = pickle.dumps(est.available_func)
+    unpickled_func = pickle.loads(pickled_bytes)
+    assert unpickled_func() == return_value
