diff --git a/pretab/transformers/binning/binning.py b/pretab/transformers/binning/binning.py index d68d52b..9f1fbf7 100644 --- a/pretab/transformers/binning/binning.py +++ b/pretab/transformers/binning/binning.py @@ -3,6 +3,7 @@ import numpy as np import pandas as pd from sklearn.base import BaseEstimator, TransformerMixin +from sklearn.utils.validation import check_is_fitted from ...core.exceptions import InsufficientSamplesError, InvalidParamError, PretabDataError from ...core.params import UNSET, AliasResolverMixin @@ -144,14 +145,17 @@ def get_feature_names_out(self, input_features=None): Parameters ---------- - input_features : list of str - The names of the input features. + input_features : list of str or None + The names of the input features. When ``None``, names of the form + ``x0, x1, ...`` are generated, matching the rest of the package and + scikit-learn's contract that the call works with no arguments. Returns ------- - input_features : ndarray of shape (n_features,) + feature_names : ndarray of shape (n_features,) The names of the output features after transformation. """ + check_is_fitted(self, "n_features_in_") if input_features is None: - raise InvalidParamError("input_features must be specified") - return input_features + input_features = [f"x{i}" for i in range(self.n_features_in_)] + return np.asarray(input_features, dtype=object) diff --git a/tests/test_custombin_transformer.py b/tests/test_custombin_transformer.py index 1be1ec5..30d19a2 100644 --- a/tests/test_custombin_transformer.py +++ b/tests/test_custombin_transformer.py @@ -69,13 +69,47 @@ def test_custom_bin_transformer_feature_names_out(): transformer = CustomBinTransformer(output_dim=3) transformer.fit(np.array([[0.2]])) names = transformer.get_feature_names_out(["feature1"]) - assert names == ["feature1"] + assert isinstance(names, np.ndarray) + assert list(names) == ["feature1"] -def test_custom_bin_transformer_feature_names_out_raises(): - transformer = CustomBinTransformer(output_dim=3) - with pytest.raises(ValueError): - transformer.get_feature_names_out() +def test_custom_bin_transformer_feature_names_out_before_fit_raises(): + from sklearn.exceptions import NotFittedError + + with pytest.raises(NotFittedError): + CustomBinTransformer(output_dim=3).get_feature_names_out() + + +# --------------------------------------------------------------------------- # +# ``get_feature_names_out()`` must work with no arguments and return an ndarray. +# +# It used to raise without ``input_features`` and return a plain list with them, +# which broke ``Pipeline.get_feature_names_out()``. The same fix was already +# applied to NoTransformer / ToFloatTransformer / ContinuousOrdinalTransformer. +# --------------------------------------------------------------------------- # +def test_feature_names_out_defaults_to_generated_names(): + transformer = CustomBinTransformer(output_dim=3).fit(np.linspace(0, 1, 10).reshape(-1, 1)) + + names = transformer.get_feature_names_out() + + assert isinstance(names, np.ndarray) + assert list(names) == ["x0"] + + +def test_feature_names_out_works_inside_a_pipeline(): + from sklearn.pipeline import Pipeline + + X = np.linspace(0, 1, 20).reshape(-1, 1) + pipe = Pipeline([("bin", CustomBinTransformer(output_dim=4))]).fit(X) + + assert list(np.asarray(pipe.get_feature_names_out())) == ["x0"] + + +def test_feature_names_out_length_matches_transform_width(): + transformer = CustomBinTransformer(output_dim=4).fit(np.linspace(0, 1, 20).reshape(-1, 1)) + + width = transformer.transform(np.linspace(0, 1, 20).reshape(-1, 1)).shape[1] + assert len(transformer.get_feature_names_out()) == width def test_custom_bin_transformer_is_sklearn_compatible():