TST add common test for string and nan (#804)

Author: glemaitre

imblearn/utils/estimator_checks.py

@@ -28,6 +28,7 @@
 from sklearn.utils.estimator_checks import _maybe_mark_xfail
 from sklearn.utils.estimator_checks import _get_check_estimator_ids
 from sklearn.utils._testing import assert_allclose
+from sklearn.utils._testing import assert_array_equal
 from sklearn.utils._testing import assert_raises_regex
 from sklearn.utils.multiclass import type_of_target
 
@@ -61,6 +62,10 @@ def _yield_sampler_checks(sampler):
         yield check_samplers_sparse
     if "dataframe" in tags["X_types"]:
         yield check_samplers_pandas
+    if "string" in tags["X_types"]:
+        yield check_samplers_string
+    if tags["allow_nan"]:
+        yield check_samplers_nan
     yield check_samplers_list
     yield check_samplers_multiclass_ova
     yield check_samplers_preserve_dtype
@@ -399,6 +404,36 @@ def check_samplers_sample_indices(name, sampler_orig):
         assert not hasattr(sampler, "sample_indices_")
 
 
+def check_samplers_string(name, sampler_orig):
+    rng = np.random.RandomState(0)
+    sampler = clone(sampler_orig)
+    categories = np.array(["A", "B", "C"], dtype=object)
+    n_samples = 30
+    X = rng.randint(low=0, high=3, size=n_samples).reshape(-1, 1)
+    X = categories[X]
+    y = rng.permutation([0] * 10 + [1] * 20)
+
+    X_res, y_res = sampler.fit_resample(X, y)
+    assert X_res.dtype == object
+    assert X_res.shape[0] == y_res.shape[0]
+    assert_array_equal(np.unique(X_res.ravel()), categories)
+
+
+def check_samplers_nan(name, sampler_orig):
+    rng = np.random.RandomState(0)
+    sampler = clone(sampler_orig)
+    categories = np.array([0, 1, np.nan], dtype=np.float64)
+    n_samples = 100
+    X = rng.randint(low=0, high=3, size=n_samples).reshape(-1, 1)
+    X = categories[X]
+    y = rng.permutation([0] * 40 + [1] * 60)
+
+    X_res, y_res = sampler.fit_resample(X, y)
+    assert X_res.dtype == np.float64
+    assert X_res.shape[0] == y_res.shape[0]
+    assert np.any(np.isnan(X_res.ravel()))
+
+
 def check_classifier_on_multilabel_or_multioutput_targets(name, estimator_orig):
     estimator = clone(estimator_orig)
     X, y = make_multilabel_classification(n_samples=30)

imblearn/utils/tests/test_estimator_checks.py

@@ -2,16 +2,18 @@
 import numpy as np
 
 from sklearn.base import BaseEstimator
-from sklearn.utils import check_X_y
 from sklearn.utils.multiclass import check_classification_targets
 
 from imblearn.base import BaseSampler
-
+from imblearn.over_sampling.base import BaseOverSampler
+from imblearn.utils import check_target_type as target_check
 from imblearn.utils.estimator_checks import check_target_type
 from imblearn.utils.estimator_checks import check_samplers_one_label
 from imblearn.utils.estimator_checks import check_samplers_fit
 from imblearn.utils.estimator_checks import check_samplers_sparse
 from imblearn.utils.estimator_checks import check_samplers_preserve_dtype
+from imblearn.utils.estimator_checks import check_samplers_string
+from imblearn.utils.estimator_checks import check_samplers_nan
 
 
 class BaseBadSampler(BaseEstimator):
@@ -64,6 +66,34 @@ def _fit_resample(self, X, y):
         return X.astype(np.float64), y.astype(np.int64)
 
 
+class IndicesSampler(BaseOverSampler):
+    def _check_X_y(self, X, y):
+        y, binarize_y = target_check(y, indicate_one_vs_all=True)
+        X, y = self._validate_data(
+            X,
+            y,
+            reset=True,
+            dtype=None,
+            force_all_finite=False,
+        )
+        return X, y, binarize_y
+
+    def _fit_resample(self, X, y):
+        n_max_count_class = np.bincount(y).max()
+        indices = np.random.choice(np.arange(X.shape[0]), size=n_max_count_class * 2)
+        return X[indices], y[indices]
+
+
+def test_check_samplers_string():
+    sampler = IndicesSampler()
+    check_samplers_string(sampler.__class__.__name__, sampler)
+
+
+def test_check_samplers_nan():
+    sampler = IndicesSampler()
+    check_samplers_nan(sampler.__class__.__name__, sampler)
+
+
 mapping_estimator_error = {
     "BaseBadSampler": (AssertionError, "ValueError not raised by fit"),
     "SamplerSingleClass": (AssertionError, "Sampler can't balance when only"),