iter

Author: glemaitre

imblearn/utils/estimator_checks.py

@@ -25,7 +25,6 @@
 from sklearn.utils.testing import assert_allclose
 from sklearn.utils.testing import set_random_state
 
-from imblearn.base import SamplerMixin
 from imblearn.over_sampling.base import BaseOverSampler
 from imblearn.under_sampling.base import BaseCleaningSampler, BaseUnderSampler
 from imblearn.ensemble.base import BaseEnsembleSampler
@@ -34,8 +33,6 @@
 
 from imblearn.utils.testing import warns
 
-NOT_TESTED_SAMPLERS = ('FunctionSampler')  # This is not a normal sampler
-
 
 def _yield_sampler_checks(name, Estimator):
     yield check_target_type
@@ -49,10 +46,10 @@ def _yield_sampler_checks(name, Estimator):
     yield check_samplers_pandas
 
 
-def _yield_all_checks(name, Estimator):
+def _yield_all_checks(name, estimator):
     # trigger our checks if this is a SamplerMixin
-    if issubclass(Estimator, SamplerMixin):
-        for check in _yield_sampler_checks(name, Estimator):
+    if hasattr(estimator, 'sample'):
+        for check in _yield_sampler_checks(name, estimator):
             yield check
 
 
@@ -75,7 +72,6 @@ def check_estimator(Estimator):
     sklearn_check_estimator(Estimator)
     check_parameters_default_constructible(name, Estimator)
     for check in _yield_all_checks(name, Estimator):
-        # if name not in NOT_TESTED_SAMPLERS:
         check(name, Estimator)