test_metaestimators.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. import pickle
  2. import pytest
  3. from sklearn.utils.metaestimators import available_if
  4. class AvailableParameterEstimator:
  5. """This estimator's `available` parameter toggles the presence of a method"""
  6. def __init__(self, available=True, return_value=1):
  7. self.available = available
  8. self.return_value = return_value
  9. @available_if(lambda est: est.available)
  10. def available_func(self):
  11. """This is a mock available_if function"""
  12. return self.return_value
  13. def test_available_if_docstring():
  14. assert "This is a mock available_if function" in str(
  15. AvailableParameterEstimator.__dict__["available_func"].__doc__
  16. )
  17. assert "This is a mock available_if function" in str(
  18. AvailableParameterEstimator.available_func.__doc__
  19. )
  20. assert "This is a mock available_if function" in str(
  21. AvailableParameterEstimator().available_func.__doc__
  22. )
  23. def test_available_if():
  24. assert hasattr(AvailableParameterEstimator(), "available_func")
  25. assert not hasattr(AvailableParameterEstimator(available=False), "available_func")
  26. def test_available_if_unbound_method():
  27. # This is a non regression test for:
  28. # https://github.com/scikit-learn/scikit-learn/issues/20614
  29. # to make sure that decorated functions can be used as an unbound method,
  30. # for instance when monkeypatching.
  31. est = AvailableParameterEstimator()
  32. AvailableParameterEstimator.available_func(est)
  33. est = AvailableParameterEstimator(available=False)
  34. with pytest.raises(
  35. AttributeError,
  36. match="This 'AvailableParameterEstimator' has no attribute 'available_func'",
  37. ):
  38. AvailableParameterEstimator.available_func(est)
  39. def test_available_if_methods_can_be_pickled():
  40. """Check that available_if methods can be pickled.
  41. Non-regression test for #21344.
  42. """
  43. return_value = 10
  44. est = AvailableParameterEstimator(available=True, return_value=return_value)
  45. pickled_bytes = pickle.dumps(est.available_func)
  46. unpickled_func = pickle.loads(pickled_bytes)
  47. assert unpickled_func() == return_value